diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 9daf59d..7fe24e2 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -1,9 +1,9 @@ name: Prek on: push: - branches: main + branches: [main] pull_request: - branches: "*" + workflow_dispatch: jobs: prek: diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index c676044..21b11ec 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -2,9 +2,9 @@ name: pytest on: push: - branches: "main" + branches: [main] pull_request: - branches: "*" + workflow_dispatch: jobs: @@ -15,4 +15,4 @@ jobs: - uses: astral-sh/setup-uv@v5 - run: uv python install - run: uv sync - - run: uv run pytest -vv + - run: uv run pytest -vv -m "not gpu and not benchmark" diff --git a/.gitignore b/.gitignore index dda9276..6286689 100644 --- a/.gitignore +++ b/.gitignore @@ -173,3 +173,8 @@ cython_debug/ # PyPI configuration file .pypirc .worktrees + +# Local-only / generated artifacts +benchmarks/ +k8s/ +unsloth_compiled_cache/ diff --git a/README.md b/README.md index e8a0681..47530d7 100644 --- a/README.md +++ b/README.md @@ -1,17 +1,19 @@ # bakery -Prompt baking distills a system prompt into model weights via KL divergence training with LoRA, so you get the behavior of a prompted model at zero inference-time prompt cost. +*Where LLMs go to get baked.* -Based on [Prompt Baking](https://arxiv.org/abs/2409.13697). +Bakery distills arbitrary **prefix contexts** — system prompts, few-shot examples, conversation histories, accumulated memories — into model weights via KL-divergence training with LoRA. You get the behavior of a context-conditioned model at zero inference-time prompt cost. + +Based on [Prompt Baking](https://arxiv.org/abs/2409.13697), generalized to arbitrary prefix contexts. ## How it works A single model serves as both teacher and student through PEFT adapter toggling: -- **Teacher** (adapters disabled): sees the system prompt, generates reference behavior -- **Student** (adapters enabled): no system prompt, trained to match the teacher's output distribution +- **Teacher** (adapters disabled): sees the full prefix context, generates reference behavior +- **Student** (adapters enabled): sees no prefix (or only the last N messages of it), trained to match the teacher's output distribution -The training objective minimizes per-token KL divergence between teacher and student logits on the response portion of each conversation. +The training objective minimizes per-token KL divergence between teacher and student logits on whichever message tokens you mark as targets (by default: all assistant turns). ## Data sources @@ -79,23 +81,35 @@ bakery --config examples/basic.yaml --num_train_epochs 5 --learning_rate 5e-5 ## As a library ```python -from bakery import BakeryConfig, PromptBakingTrainer, PromptBakingDataset, prompt_baking_collator +from bakery import ( + BakeryConfig, + ContextConfig, + ContextBakingTrainer, + create_dataset, + prompt_baking_collator, +) config = BakeryConfig( output_dir="./outputs", - system_prompt="You are helpful.", num_train_epochs=3, learning_rate=1e-4, ) -dataset = PromptBakingDataset( - prompts=["What is AI?", "Explain gravity."], - responses=["AI is...", "Gravity is..."], # optional precomputed +context_config = ContextConfig( + prefix_messages=[ + {"role": "system", "content": "You are helpful."}, + ], +) + +dataset = create_dataset( + ["What is AI?", "Explain gravity."], + ["AI is...", "Gravity is..."], # optional precomputed responses ) -trainer = PromptBakingTrainer( +trainer = ContextBakingTrainer( model=peft_model, args=config, + context_config=context_config, train_dataset=dataset, processing_class=tokenizer, data_collator=prompt_baking_collator, @@ -103,9 +117,48 @@ trainer = PromptBakingTrainer( trainer.train() ``` +## Context baking + +Beyond a single system prompt, bakery supports arbitrary prefix contexts via `ContextConfig` — flat YAML fields alongside the rest: + +```yaml +# Any list of chat messages. Teacher sees these prepended to every example. +prefix_messages: + - {role: system, content: "You answer concisely."} + - {role: user, content: "Example Q"} + - {role: assistant, content: "Example A"} + +# How many trailing prefix messages the *student* also sees. +# 0 (default) = pure baking. N>0 = last N messages are kept at inference. +student_retained_turns: 0 + +# Which message roles contribute tokens to the KL loss. +target_roles: [assistant] + +# Optional regex (re.search) over message content to further restrict targets. +# Useful to bake only final answers while ignoring chain-of-thought. +target_content_pattern: "^Answer:" +``` + +Alternatively, load the prefix from a file: + +```yaml +prefix_messages_file: "./prefixes/persona_A.yaml" # or .json +``` + +**Per-row prefixes.** A `prefix_messages` column on the dataset overrides the global one, so you can bake many contexts into a single adapter. + +**Multi-turn datasets.** HuggingFace chat datasets (`messages` column) are loaded verbatim: each row's full conversation becomes the teacher view, and `student_retained_turns` controls how much of it the student sees. + +**Migration from `system_prompt`.** The old `system_prompt: "..."` field still works — it auto-desugars to `prefix_messages: [{role: system, content: ...}]` and emits a `DeprecationWarning`. Prefer the new field for new configs. + ## Examples | Config | Description | |--------|-------------| | [`examples/basic.yaml`](examples/basic.yaml) | On-the-fly trajectory generation from inline prompts | | [`examples/sft_dataset.yaml`](examples/sft_dataset.yaml) | Bake from an existing HF chat dataset | +| [`examples/multi_turn_prefix.yaml`](examples/multi_turn_prefix.yaml) | System prompt + few-shot demonstration prefix | +| [`examples/continual_memory.yaml`](examples/continual_memory.yaml) | Multi-turn HF chat dataset with `student_retained_turns: 2` | +| [`examples/per_row_prefix.yaml`](examples/per_row_prefix.yaml) | Per-row persona/context prefixes from a JSON dataset | +| [`examples/pattern_targets.yaml`](examples/pattern_targets.yaml) | Regex-filtered KL targets (bake final answers only) | diff --git a/examples/continual_memory.yaml b/examples/continual_memory.yaml new file mode 100644 index 0000000..a4d3368 --- /dev/null +++ b/examples/continual_memory.yaml @@ -0,0 +1,50 @@ +# Continual-memory baking from a multi-turn chat dataset. +# +# Each row of the dataset is a full conversation. The teacher sees the +# whole history; the student sees only the last 2 messages (e.g. the most +# recent user turn + its target response). KL distills the benefit of +# long-range memory into LoRA, so the student behaves as if it still +# remembered everything. +# +# Usage: +# bakery --config examples/continual_memory.yaml + +# --- Standard TrainingArguments --- +output_dir: "./outputs/continual_memory" +num_train_epochs: 1 +learning_rate: 5e-5 +per_device_train_batch_size: 1 +gradient_accumulation_steps: 8 +logging_steps: 10 +save_strategy: "epoch" +bf16: true +report_to: "none" +seed: 42 +max_seq_length: 2048 + +# --- Context baking --- +# No global prefix — per-row prefix comes from the dataset (multi-turn +# `messages` column). Student retains the last 2 messages of each row. +student_retained_turns: 2 +target_roles: + - assistant + +# --- Model --- +model_name_or_path: "Qwen/Qwen3-0.6B" +torch_dtype: "bfloat16" +attn_implementation: "sdpa" + +# --- LoRA --- +r: 32 +lora_alpha: 64 +target_modules: + - q_proj + - k_proj + - v_proj + - o_proj +lora_dropout: 0.05 + +# --- Data: HuggingFace multi-turn chat dataset --- +# The loader reads the `messages` column and preserves all turns. +dataset: "HuggingFaceH4/ultrachat_200k" +dataset_split: "train_sft[:500]" diff --git a/examples/multi_turn_prefix.yaml b/examples/multi_turn_prefix.yaml new file mode 100644 index 0000000..663817c --- /dev/null +++ b/examples/multi_turn_prefix.yaml @@ -0,0 +1,62 @@ +# Multi-turn prefix baking: system + few-shot examples. +# +# The teacher sees a system prompt plus a few-shot user/assistant pair +# before each training example; the student sees none of that (pure +# baking). KL distills the conditioned-teacher behavior into LoRA. +# +# Usage: +# bakery --config examples/multi_turn_prefix.yaml + +# --- Standard TrainingArguments --- +output_dir: "./outputs/multi_turn_prefix" +num_train_epochs: 3 +learning_rate: 1e-4 +per_device_train_batch_size: 1 +gradient_accumulation_steps: 4 +logging_steps: 5 +save_strategy: "epoch" +bf16: true +report_to: "none" +seed: 42 + +# --- Context baking (prefix = system + few-shot) --- +prefix_messages: + - role: system + content: "You answer concisely, in one short sentence." + - role: user + content: "What is the capital of France?" + - role: assistant + content: "Paris." + - role: user + content: "What color is the sky?" + - role: assistant + content: "Blue." + +# Pure baking: student sees no prefix at all. +student_retained_turns: 0 + +# Only assistant tokens contribute to KL loss (default). +target_roles: + - assistant + +# --- Model --- +model_name_or_path: "Qwen/Qwen3-0.6B" +torch_dtype: "bfloat16" +attn_implementation: "sdpa" + +# --- LoRA --- +r: 32 +lora_alpha: 64 +target_modules: + - q_proj + - k_proj + - v_proj + - o_proj +lora_dropout: 0.05 + +# --- Data --- +training_prompts: + - "How many legs does a spider have?" + - "What is the boiling point of water?" + - "Name a prime number greater than ten." + - "Who wrote Hamlet?" diff --git a/examples/pattern_targets.yaml b/examples/pattern_targets.yaml new file mode 100644 index 0000000..27c06aa --- /dev/null +++ b/examples/pattern_targets.yaml @@ -0,0 +1,62 @@ +# Pattern-filtered target masking. +# +# Not every assistant turn deserves KL loss — only those whose content +# matches a regex. Useful when your data mixes chain-of-thought with +# final answers and you only want to distill the final answers, or when +# you want to bake only the "structured" portion of multi-style replies. +# +# Here: only assistant messages whose content starts with "Answer:" are +# treated as targets. Everything else (system, user, non-matching +# assistant turns) is passed through to both teacher and student but +# contributes zero loss. +# +# Usage: +# bakery --config examples/pattern_targets.yaml + +# --- Standard TrainingArguments --- +output_dir: "./outputs/pattern_targets" +num_train_epochs: 2 +learning_rate: 1e-4 +per_device_train_batch_size: 2 +gradient_accumulation_steps: 2 +logging_steps: 5 +save_strategy: "epoch" +bf16: true +report_to: "none" +seed: 42 + +# --- Context baking --- +prefix_messages: + - role: system + content: >- + Reason step-by-step, then give a final answer on a new line + beginning with "Answer: ". + +student_retained_turns: 0 +target_roles: + - assistant + +# Only tokens from assistant messages that match this regex (re.search) +# are scored by KL. CoT / reasoning prefixes are ignored. +target_content_pattern: "^Answer:" + +# --- Model --- +model_name_or_path: "Qwen/Qwen3-0.6B" +torch_dtype: "bfloat16" +attn_implementation: "sdpa" + +# --- LoRA --- +r: 32 +lora_alpha: 64 +target_modules: + - q_proj + - k_proj + - v_proj + - o_proj +lora_dropout: 0.05 + +# --- Data --- +training_prompts: + - "What is 17 times 3?" + - "If a train leaves at 3pm traveling 60mph, where is it at 5pm?" + - "How many vowels are in the word 'strawberry'?" diff --git a/examples/per_row_prefix.yaml b/examples/per_row_prefix.yaml new file mode 100644 index 0000000..5b43760 --- /dev/null +++ b/examples/per_row_prefix.yaml @@ -0,0 +1,63 @@ +# Per-row prefix baking (e.g. many personas, one model). +# +# Each training row carries its own `prefix_messages` column that +# overrides the global one. This is useful for baking a family of +# personas/contexts into a single adapter — each example conditions the +# teacher on a different prefix and the student learns them all. +# +# Data format (JSON): +# [ +# { +# "prefix_messages": [{"role": "system", "content": "Persona A ..."}], +# "messages": [ +# {"role": "user", "content": "..."}, +# {"role": "assistant", "content": "..."} +# ] +# }, +# ... +# ] +# +# Usage: +# bakery --config examples/per_row_prefix.yaml + +# --- Standard TrainingArguments --- +output_dir: "./outputs/per_row_prefix" +num_train_epochs: 2 +learning_rate: 1e-4 +per_device_train_batch_size: 1 +gradient_accumulation_steps: 4 +logging_steps: 5 +save_strategy: "epoch" +bf16: true +report_to: "none" +seed: 42 + +# --- Context baking --- +# Fallback prefix used if a row has no `prefix_messages`. Individual rows +# override this via their own `prefix_messages` field. +prefix_messages: + - role: system + content: "You are a helpful assistant." + +student_retained_turns: 0 +target_roles: + - assistant + +# --- Model --- +model_name_or_path: "Qwen/Qwen3-0.6B" +torch_dtype: "bfloat16" +attn_implementation: "sdpa" + +# --- LoRA --- +r: 32 +lora_alpha: 64 +target_modules: + - q_proj + - k_proj + - v_proj + - o_proj +lora_dropout: 0.05 + +# --- Data --- +# Local JSON file with per-row prefix_messages + messages. +dataset: "./data/personas.json" diff --git a/src/bakery/__init__.py b/src/bakery/__init__.py index a9b5830..066bcd2 100644 --- a/src/bakery/__init__.py +++ b/src/bakery/__init__.py @@ -1,20 +1,32 @@ """Bakery - Where LLMs go to get baked. -Prompt baking via KL divergence distillation with LoRA. +Context baking (prefix-context distillation) via KL divergence with LoRA. """ -from bakery.config import BakeryConfig, DataConfig, LoraConfig -from bakery.trainer import PromptBakingTrainer -from bakery.data import create_dataset, prompt_baking_collator, load_dataset +from bakery.config import BakeryConfig, ContextConfig, DataConfig, LoraConfig +from bakery.trainer import ContextBakingTrainer, PromptBakingTrainer +from bakery.data import ( + create_conversational_dataset, + create_dataset, + load_conversations, + load_dataset, + prompt_baking_collator, +) from bakery.kl import compute_kl_divergence +from bakery.masking import build_target_mask __all__ = [ "BakeryConfig", + "ContextConfig", "DataConfig", "LoraConfig", + "ContextBakingTrainer", "PromptBakingTrainer", + "create_conversational_dataset", "create_dataset", - "prompt_baking_collator", + "load_conversations", "load_dataset", + "prompt_baking_collator", "compute_kl_divergence", + "build_target_mask", ] diff --git a/src/bakery/cli.py b/src/bakery/cli.py index a763ac3..06fefb6 100644 --- a/src/bakery/cli.py +++ b/src/bakery/cli.py @@ -8,6 +8,7 @@ import argparse import os import json +import warnings import torch from transformers import ( @@ -18,11 +19,13 @@ ) from peft import LoraConfig as PeftLoraConfig, get_peft_model -from bakery.config import BakeryConfig, DataConfig, LoraConfig -from bakery.trainer import PromptBakingTrainer +from bakery.config import BakeryConfig, ContextConfig, DataConfig, LoraConfig +from bakery.trainer import ContextBakingTrainer from bakery.data import ( + create_conversational_dataset, create_dataset, prompt_baking_collator, + load_conversations, load_corpus, build_system_prompt, load_data, @@ -37,6 +40,24 @@ } +def _load_prefix_file(path: str) -> list: + """Load prefix_messages from a JSON or YAML file (expects a list of {role, content}).""" + with open(path) as f: + text = f.read() + if path.endswith((".yaml", ".yml")): + import yaml + + data = yaml.safe_load(text) + else: + data = json.loads(text) + if not isinstance(data, list): + raise ValueError( + f"prefix_messages_file {path!r} must contain a JSON/YAML list of " + "{role, content} dicts." + ) + return data + + def main(): # Pre-parse --config and collect remaining args for HfArgumentParser. # argparse handles -h/--help and validates --config before we proceed. @@ -57,9 +78,9 @@ def main(): pre_args, remaining_args = pre_parser.parse_known_args() config_file = pre_args.config - parser = HfArgumentParser((BakeryConfig, DataConfig, LoraConfig)) + parser = HfArgumentParser((BakeryConfig, DataConfig, LoraConfig, ContextConfig)) - baking_config, data_config, lora_config = parser.parse_yaml_file( + baking_config, data_config, lora_config, context_config = parser.parse_yaml_file( config_file, allow_extra_keys=True ) # Apply CLI overrides on top of YAML config. @@ -73,38 +94,96 @@ def main(): if arg.startswith("--"): explicit_keys.add(arg.lstrip("-").replace("-", "_")) - override_parser = HfArgumentParser((BakeryConfig, DataConfig, LoraConfig)) + override_parser = HfArgumentParser( + (BakeryConfig, DataConfig, LoraConfig, ContextConfig) + ) overrides = override_parser.parse_args_into_dataclasses( args=["--output_dir", baking_config.output_dir] + remaining_args, return_remaining_strings=True, ) for override_cfg, base_cfg in zip( - overrides[:3], - (baking_config, data_config, lora_config), + overrides[:4], + (baking_config, data_config, lora_config, context_config), ): for k, v in vars(override_cfg).items(): if k in explicit_keys: setattr(base_cfg, k, v) - # Build system prompt + # Load prefix_messages_file if set (JSON or YAML list of {role, content} dicts). + if context_config.prefix_messages is None and context_config.prefix_messages_file: + context_config.prefix_messages = _load_prefix_file( + context_config.prefix_messages_file + ) + + # Build system prompt (supports corpus-based knowledge baking) — retained + # for backward compat; desugars into prefix_messages below. corpus = load_corpus(data_config) - system_prompt = build_system_prompt(baking_config, data_config, corpus) - baking_config.system_prompt = system_prompt + try: + system_prompt = build_system_prompt(baking_config, data_config, corpus) + except ValueError: + # No system_prompt configured — OK if prefix_messages is set. + system_prompt = None + if system_prompt: + baking_config.system_prompt = system_prompt + + # Desugar deprecated system_prompt → prefix_messages when user didn't set prefix. + if not context_config.prefix_messages and baking_config.system_prompt: + if baking_config.system_prompt_file or corpus: + pass # typical (non-deprecated) path: corpus-driven or file-loaded prompt + else: + warnings.warn( + "`system_prompt` is deprecated; use ContextConfig.prefix_messages instead. " + "Auto-wrapping into prefix_messages=[{role: system, content: ...}].", + DeprecationWarning, + stacklevel=2, + ) + context_config.prefix_messages = [ + {"role": "system", "content": baking_config.system_prompt} + ] + + if not context_config.prefix_messages: + raise ValueError( + "No prefix context configured. Set prefix_messages (inline or via " + "prefix_messages_file), or the deprecated system_prompt / corpus_file." + ) - # Load data - training_prompts, precomputed_responses = load_data(data_config) + # Load data. Use conversational loader when the source preserves multi-turn + # history (HF `messages` column or JSON rows with `messages`/`prefix_messages`); + # otherwise use the simple (prompts, responses) path. + conversational_rows = None + training_prompts, precomputed_responses = [], None + if data_config.dataset or data_config.training_prompts: + if data_config.dataset: + try: + conversational_rows = load_conversations( + data_config.dataset, data_config.dataset_split + ) + except Exception: + conversational_rows = None + if conversational_rows is None: + training_prompts, precomputed_responses = load_data(data_config) eval_qa = load_eval_data(data_config.eval_file) heldout_qa = load_eval_data(data_config.heldout_file) print("=" * 70) - print("Bakery - Prompt Baking with KL Divergence") + print("Bakery - Context Baking with KL Divergence") print("=" * 70) - print(f" System prompt: {len(system_prompt):,} chars") - print(f" Training prompts: {len(training_prompts)}") - if precomputed_responses: - print(f" Precomputed responses: {len(precomputed_responses)}") + print( + f" Prefix messages: {len(context_config.prefix_messages)} " + f"({sum(len(m.get('content', '')) for m in context_config.prefix_messages):,} chars)" + ) + if conversational_rows is not None: + n_turns = sum(len(r.get("turns", [])) for r in conversational_rows) + print( + f" Training rows (conversational): {len(conversational_rows)} " + f"({n_turns} total turns)" + ) else: - print(" Mode: on-the-fly trajectory generation") + print(f" Training prompts: {len(training_prompts)}") + if precomputed_responses: + print(f" Precomputed responses: {len(precomputed_responses)}") + else: + print(" Mode: on-the-fly trajectory generation") if eval_qa: print(f" Evaluation Q&A: {len(eval_qa)}") if heldout_qa: @@ -127,6 +206,37 @@ def main(): features.append("unsloth") ensure_deps(model_type=model_type, features=features) + # Gemma 4 introduces Gemma4ClippableLinear (inherits nn.Module, not nn.Linear) + # in its vision/audio encoders. PEFT rejects it even when targeting only text + # layers. Monkey-patch it to inherit nn.Linear so PEFT's type check passes. + # Must happen before from_pretrained() materialises the model. + try: + from transformers.models.gemma4 import modeling_gemma4 + import torch.nn as nn + + class _PatchedClippableLinear(nn.Linear): + def __init__(self, config, in_features, out_features): + nn.Linear.__init__(self, in_features, out_features, bias=False) + self.use_clipped_linears = getattr(config, "use_clipped_linears", False) + if self.use_clipped_linears: + self.register_buffer("input_min", torch.tensor(-float("inf"))) + self.register_buffer("input_max", torch.tensor(float("inf"))) + self.register_buffer("output_min", torch.tensor(-float("inf"))) + self.register_buffer("output_max", torch.tensor(float("inf"))) + + def forward(self, x): + if self.use_clipped_linears: + x = torch.clamp(x, self.input_min, self.input_max) + out = nn.Linear.forward(self, x) + if self.use_clipped_linears: + out = torch.clamp(out, self.output_min, self.output_max) + return out + + modeling_gemma4.Gemma4ClippableLinear = _PatchedClippableLinear + print(" Patched Gemma4ClippableLinear for PEFT compatibility") + except (ImportError, AttributeError): + pass # Not a Gemma 4 run or transformers too old + # Load model print(f"\n[1] Loading model: {data_config.model_name_or_path}") torch_dtype = DTYPE_MAP.get(data_config.torch_dtype, torch.bfloat16) @@ -209,32 +319,47 @@ def main(): if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - prompt_tokens = len(tokenizer.encode(system_prompt)) - print(f" System prompt tokens: {prompt_tokens:,}") - - # Baseline evaluations + # Rough prefix-token count (for summary output only). + prefix_rendered = tokenizer.apply_chat_template( + context_config.prefix_messages, tokenize=False, add_generation_prompt=False + ) + prompt_tokens = len(tokenizer.encode(prefix_rendered, add_special_tokens=False)) + print(f" Prefix tokens: {prompt_tokens:,}") + + # Baseline evaluations (only when prefix is a single system message; otherwise + # "discrete prompt" baseline comparison isn't well-defined). + _is_simple_system_prompt = ( + len(context_config.prefix_messages) == 1 + and context_config.prefix_messages[0].get("role") == "system" + ) + discrete_sp = ( + context_config.prefix_messages[0]["content"] + if _is_simple_system_prompt + else None + ) if eval_qa: print("\n[2] Baseline evaluations...") - print("\n === No system prompt ===") + print("\n === No prefix ===") baseline_none = evaluate_model(base_model, tokenizer, eval_qa, "Baseline") - print("\n === With system prompt (discrete) ===") - baseline_discrete = evaluate_model( - base_model, - tokenizer, - eval_qa, - "Discrete prompt", - system_prompt=system_prompt, - ) - heldout_discrete = None - if heldout_qa: - print("\n === Held-out (discrete prompt) ===") - heldout_discrete = evaluate_model( + baseline_discrete = heldout_discrete = None + if discrete_sp: + print("\n === With system prompt (discrete) ===") + baseline_discrete = evaluate_model( base_model, tokenizer, - heldout_qa, - "Discrete (held-out)", - system_prompt=system_prompt, + eval_qa, + "Discrete prompt", + system_prompt=discrete_sp, ) + if heldout_qa: + print("\n === Held-out (discrete prompt) ===") + heldout_discrete = evaluate_model( + base_model, + tokenizer, + heldout_qa, + "Discrete (held-out)", + system_prompt=discrete_sp, + ) else: baseline_none = baseline_discrete = heldout_discrete = None @@ -272,16 +397,19 @@ def main(): # Create dataset print("\n[4] Training with KL divergence...") - if precomputed_responses: - print( - f" Using {len(precomputed_responses)} precomputed (prompt, response) pairs" - ) + if conversational_rows is not None: + train_dataset = create_conversational_dataset(conversational_rows) + print(f" Using {len(train_dataset)} conversational rows") else: - print( - f" Generating {baking_config.num_trajectories} trajectories per prompt on-the-fly" - ) - - train_dataset = create_dataset(training_prompts, precomputed_responses) + if precomputed_responses: + print( + f" Using {len(precomputed_responses)} precomputed (prompt, response) pairs" + ) + else: + print( + f" Generating {baking_config.num_trajectories} trajectories per prompt on-the-fly" + ) + train_dataset = create_dataset(training_prompts, precomputed_responses) eval_dataset = None _eval_source = data_config.eval_dataset or data_config.dataset @@ -291,23 +419,28 @@ def main(): f" Loading eval split: {_eval_split}" + (f" from {data_config.eval_dataset}" if data_config.eval_dataset else "") ) - eval_prompts, eval_responses = load_data( - type( - "_DC", - (), - { - "dataset": _eval_source, - "dataset_split": _eval_split, - "training_prompts": None, - }, - )() - ) - eval_dataset = create_dataset(eval_prompts, eval_responses) + try: + eval_rows = load_conversations(_eval_source, _eval_split) + eval_dataset = create_conversational_dataset(eval_rows) + except Exception: + eval_prompts, eval_responses = load_data( + type( + "_DC", + (), + { + "dataset": _eval_source, + "dataset_split": _eval_split, + "training_prompts": None, + }, + )() + ) + eval_dataset = create_dataset(eval_prompts, eval_responses) print(f" Eval samples: {len(eval_dataset)}") - trainer = PromptBakingTrainer( + trainer = ContextBakingTrainer( model=model, args=baking_config, + context_config=context_config, train_dataset=train_dataset, eval_dataset=eval_dataset, processing_class=tokenizer, @@ -354,7 +487,7 @@ def main(): print("\n" + "=" * 70) print("SUMMARY") print("=" * 70) - print(f"System prompt: {prompt_tokens:,} tokens -> 0 tokens at inference") + print(f"Prefix context: {prompt_tokens:,} tokens -> 0 tokens at inference") if eval_qa and baseline_discrete and baked_eval: print("\nEvaluation set:") diff --git a/src/bakery/config.py b/src/bakery/config.py index 2fd0ffa..6233115 100644 --- a/src/bakery/config.py +++ b/src/bakery/config.py @@ -23,11 +23,17 @@ class BakeryConfig(TrainingArguments): system_prompt: Optional[str] = field( default=None, - metadata={"help": "System prompt text to bake into model weights."}, + metadata={ + "help": "DEPRECATED — use ContextConfig.prefix_messages instead. " + "When set, desugars to prefix_messages=[{role: system, content: ...}]." + }, ) system_prompt_file: Optional[str] = field( default=None, - metadata={"help": "Path to file containing the system prompt."}, + metadata={ + "help": "DEPRECATED — use ContextConfig.prefix_messages_file instead. " + "Path to file containing the system prompt." + }, ) num_trajectories: int = field( default=4, @@ -240,3 +246,54 @@ def __post_init__(self): lora_dropout: float = field(default=0.05, metadata={"help": "LoRA dropout."}) bias: str = field(default="none", metadata={"help": "LoRA bias mode."}) + + +@dataclass +class ContextConfig: + """Prefix context and target-mask configuration for context baking. + + Generalizes system-prompt baking to arbitrary prefix contexts (conversation + histories, accumulated memories, few-shot examples). The teacher sees the + full prefix; the student sees an optionally-trimmed version. KL is computed + only on tokens matching `target_roles` / `target_content_pattern`. + """ + + prefix_messages: Optional[List[dict]] = field( + default=None, + metadata={ + "help": "Global prefix context as a list of {role, content} dicts. " + "Teacher sees this prepended to every example; student does not " + "(or sees only the last student_retained_turns of it). Overridden " + "per-row by a 'prefix_messages' dataset column if present." + }, + ) + prefix_messages_file: Optional[str] = field( + default=None, + metadata={ + "help": "Path to a JSON or YAML file containing the prefix_messages list. " + "Loaded at CLI parse time if prefix_messages is not set inline." + }, + ) + student_retained_turns: int = field( + default=0, + metadata={ + "help": "Number of trailing prefix messages the student also sees. " + "0 (default) = pure baking: student sees no prefix. N>0 = student " + "sees the last N messages of the prefix; earlier messages are baked." + }, + ) + target_roles: List[str] = field( + default_factory=lambda: ["assistant"], + metadata={ + "help": "Message roles whose tokens receive KL loss. Default: ['assistant']. " + "Any message whose role is in this list is treated as a training target." + }, + ) + target_content_pattern: Optional[str] = field( + default=None, + metadata={ + "help": "Optional regex applied to message content. When set, a message " + "becomes a target only if its role is in target_roles AND its content " + "matches this pattern (re.search semantics)." + }, + ) diff --git a/src/bakery/data.py b/src/bakery/data.py index a3218a4..c2f3146 100644 --- a/src/bakery/data.py +++ b/src/bakery/data.py @@ -1,7 +1,7 @@ -"""Data loading utilities for prompt baking.""" +"""Data loading utilities for context baking.""" import json -from typing import List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple from datasets import Dataset @@ -12,7 +12,7 @@ def create_dataset( prompts: List[str], responses: Optional[List[str]] = None, ) -> Dataset: - """Create a HuggingFace Dataset for prompt baking training. + """Create a HuggingFace Dataset for single-turn context baking training. Args: prompts: List of user messages. @@ -25,10 +25,52 @@ def create_dataset( return Dataset.from_dict(data) +def create_conversational_dataset( + rows: List[Dict[str, Any]], +) -> Dataset: + """Create a HuggingFace Dataset from conversational rows. + + Each row is a dict with: + - `turns`: List[{role, content}] — the training conversation (required). + - `prefix_messages`: Optional[List[{role, content}]] — per-row prefix + that overrides the global ContextConfig.prefix_messages. + - `response`: Optional[str] — for trajectory mode or single-turn data; + when set, the trainer appends {role: assistant, content: response} + to `turns`. + """ + data = { + "turns": [r.get("turns", []) for r in rows], + "prefix_messages": [r.get("prefix_messages") for r in rows], + "responses": [r.get("response") for r in rows], + } + return Dataset.from_dict(data) + + def prompt_baking_collator(features: List[dict]) -> dict: - """Collate function that passes user_messages and responses through.""" - user_messages, responses = [], [] + """Collate function for context baking. + + Accepts both shapes and emits a batch with whichever keys are present: + + - Legacy: {user_messages, responses} + - Conversational: {prefix_messages, turns, responses} + + Legacy items are passed through unchanged so the trainer can normalize + them to the conversational shape. + """ + user_messages: List[str] = [] + responses: List[Optional[str]] = [] + turns: List[List[dict]] = [] + prefix_messages: List[Optional[List[dict]]] = [] + saw_turns = False + for f in features: + if "turns" in f: + saw_turns = True + turns.append(list(f["turns"]) if f["turns"] is not None else []) + prefix_messages.append(f.get("prefix_messages")) + responses.append(f.get("responses")) + continue + msg = f.get("user_messages") if isinstance(msg, list): user_messages.extend(msg) @@ -40,6 +82,13 @@ def prompt_baking_collator(features: List[dict]) -> dict: responses.extend(resp) elif isinstance(resp, str): responses.append(resp) + + if saw_turns: + return { + "turns": turns, + "prefix_messages": prefix_messages, + "responses": responses, + } return {"user_messages": user_messages, "responses": responses} @@ -243,6 +292,135 @@ def _load_hf( return prompts, None +def load_conversations(source: str, split: str = "train") -> List[Dict[str, Any]]: + """Load conversational training rows from a local JSON file or HF dataset. + + Returns a list of dicts suitable for `create_conversational_dataset`: + - `turns`: List[{role, content}] + - `prefix_messages`: Optional[List[{role, content}]] + - `response`: Optional[str] + + HF `messages` columns are preserved as full turn sequences (NOT flattened), + so multi-turn conversations can be baked with all assistant turns as KL targets. + """ + import os + + if os.path.exists(source): + return _load_conversations_json(source) + return _load_conversations_hf(source, split) + + +def _load_conversations_json(path: str) -> List[Dict[str, Any]]: + with open(path) as f: + data = json.load(f) + + if isinstance(data, dict): + for key in ( + "pairs", + "data", + "samples", + "completions", + "training_samples", + "conversations", + ): + if key in data: + data = data[key] + break + + if isinstance(data, list) and data and isinstance(data[0], str): + return [ + { + "turns": [{"role": "user", "content": p}], + "prefix_messages": None, + "response": None, + } + for p in data + ] + + rows: List[Dict[str, Any]] = [] + for item in data: + prefix = item.get("prefix_messages") + if "messages" in item and isinstance(item["messages"], list): + turns = item["messages"] + response = None + else: + p = item.get( + "prompt", + item.get("user_message", item.get("input", item.get("question", ""))), + ) + if not p: + continue + turns = [{"role": "user", "content": p}] + response = item.get("response", item.get("completion", item.get("output"))) + rows.append( + { + "turns": list(turns), + "prefix_messages": list(prefix) if prefix else None, + "response": response, + } + ) + return rows + + +def _load_conversations_hf( + dataset_id: str, split: str = "train" +) -> List[Dict[str, Any]]: + from datasets import load_dataset as hf_load_dataset + + ds = hf_load_dataset(dataset_id, split=split) + columns = ds.column_names + has_prefix = "prefix_messages" in columns + + if "messages" in columns: + rows: List[Dict[str, Any]] = [] + for row in ds: + msgs = row["messages"] + if not msgs: + continue + rows.append( + { + "turns": list(msgs), + "prefix_messages": list(row["prefix_messages"]) + if has_prefix and row.get("prefix_messages") + else None, + "response": None, + } + ) + return rows + + prompt_col = None + for col in ("prompt", "input", "question", "text", "instruction"): + if col in columns: + prompt_col = col + break + if prompt_col is None: + raise ValueError( + f"Cannot find prompt column in dataset '{dataset_id}'. " + f"Available columns: {columns}" + ) + + response_col = None + for col in ("response", "completion", "output", "answer", "target"): + if col in columns: + response_col = col + break + + rows = [] + for row in ds: + if not row.get(prompt_col): + continue + rows.append( + { + "turns": [{"role": "user", "content": row[prompt_col]}], + "prefix_messages": list(row["prefix_messages"]) + if has_prefix and row.get("prefix_messages") + else None, + "response": row[response_col] if response_col else None, + } + ) + return rows + + def load_eval_data(eval_file: Optional[str]) -> List[Tuple[str, List[str]]]: """Load evaluation Q&A pairs from file.""" if not eval_file: diff --git a/src/bakery/evaluate.py b/src/bakery/evaluate.py index 8daa594..bfc3b9f 100644 --- a/src/bakery/evaluate.py +++ b/src/bakery/evaluate.py @@ -49,6 +49,14 @@ def evaluate_model( model.device ) + # Gemma 4 requires token_type_ids and mm_token_type_ids even for text-only. + _mtype = getattr(model.config, "model_type", None) + if _mtype in ("gemma3", "gemma4", "gemma4_text"): + if "token_type_ids" not in inputs: + inputs["token_type_ids"] = torch.zeros_like(inputs["input_ids"]) + if _mtype in ("gemma4", "gemma4_text"): + inputs["mm_token_type_ids"] = torch.zeros_like(inputs["input_ids"]) + with torch.no_grad(): outputs = model.generate( **inputs, diff --git a/src/bakery/masking.py b/src/bakery/masking.py new file mode 100644 index 0000000..ec586fe --- /dev/null +++ b/src/bakery/masking.py @@ -0,0 +1,172 @@ +"""Target-mask construction for context baking. + +Given a list of chat messages, produce a per-token boolean mask indicating +which tokens are training targets (i.e. tokens that should receive KL loss). +Selection is controlled by `target_roles` (message role filter) and an optional +`target_content_pattern` (regex on message content). + +Strategy: tokenize incremental prefixes of the message list and derive per-message +token spans by diffing cumulative token counts. For each message that matches the +target criteria, its span in the final tokenized sequence is marked True. + +This approach is tokenizer- and chat-template-agnostic. It assumes cumulative +`apply_chat_template(messages[:k], tokenize=True)` produces a token-prefix of +`apply_chat_template(messages[:k+1], tokenize=True)`. This holds for well-formed +chat templates (Llama, Qwen, Gemma, Mistral). A longest-common-token-prefix +fallback handles templates that close blocks at the end of each render. +""" + +from __future__ import annotations + +import re +from typing import List, Optional, Tuple + +import torch + + +def _tokenize_prefix(tokenizer, messages: List[dict]) -> List[int]: + """Render and tokenize a message-list prefix; return input_ids as a list.""" + if not messages: + return [] + rendered = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + ) + ids = tokenizer(rendered, add_special_tokens=False, return_tensors=None)[ + "input_ids" + ] + # Some tokenizers return List[List[int]] when given a single string; flatten. + if ids and isinstance(ids[0], list): + ids = ids[0] + return list(ids) + + +def _longest_common_prefix_len(a: List[int], b: List[int]) -> int: + """Length of the longest shared token prefix between two id sequences.""" + n = min(len(a), len(b)) + i = 0 + while i < n and a[i] == b[i]: + i += 1 + return i + + +def _per_message_spans( + tokenizer, messages: List[dict] +) -> Tuple[List[int], List[Tuple[int, int]]]: + """Return (full_input_ids, spans) where spans[i] is (start, end) of messages[i] + in the full tokenized sequence. + + Uses longest-common-token-prefix to be resilient to templates that close + blocks at each render (the closing tokens of messages[:k] won't appear in + messages[:k+1], so we take the common prefix as the "fixed" boundary). + """ + full = _tokenize_prefix(tokenizer, messages) + spans: List[Tuple[int, int]] = [] + prev_boundary = 0 + prev_tokens: List[int] = [] + for i in range(len(messages)): + this_tokens = _tokenize_prefix(tokenizer, messages[: i + 1]) + # Boundary between message i-1 and i is the longest common prefix + # between prev_tokens (messages[:i]) and this_tokens (messages[:i+1]). + if i == 0: + start = 0 + else: + start = _longest_common_prefix_len(prev_tokens, this_tokens) + # Clamp to the previous boundary: each message can only extend forward. + start = max(start, prev_boundary) + # End of message i = longest common prefix between this_tokens and full. + end = _longest_common_prefix_len(this_tokens, full) + end = max(end, start) + spans.append((start, end)) + prev_boundary = end + prev_tokens = this_tokens + return full, spans + + +def _messages_hash(messages: List[dict]) -> str: + """Stable content hash of a message list for caching.""" + import hashlib + + h = hashlib.blake2b(digest_size=16) + for m in messages: + h.update(str(m.get("role", "")).encode()) + h.update(b"\x1f") + h.update(str(m.get("content", "")).encode()) + h.update(b"\x1e") + return h.hexdigest() + + +# Bounded cache; tokenizer identity is part of the key via id(). +_MASK_CACHE: "dict[tuple, tuple[list[int], torch.BoolTensor, int]]" = {} +_MASK_CACHE_MAX = 10_000 + + +def build_target_mask( + tokenizer, + messages: List[dict], + target_roles: List[str], + content_pattern: Optional[str] = None, + target_min_msg_idx: int = 0, +) -> Tuple[List[int], torch.BoolTensor, int]: + """Compute per-token target mask for a message list. + + Args: + tokenizer: HF tokenizer with `apply_chat_template` support. + messages: list of {role, content} dicts (full sequence including targets). + target_roles: roles whose tokens should be KL targets. + content_pattern: optional regex; if set, message must also match via re.search. + target_min_msg_idx: messages at indices < this are never targets (used to + exclude the prefix portion when it happens to contain assistant turns + like few-shot examples that should be baked, not trained on). + + Returns: + (input_ids, mask, first_target_idx) where: + - input_ids: tokenized full sequence (list[int]) + - mask: BoolTensor of shape (len(input_ids),), True for target tokens + - first_target_idx: index of the first True token (or len(input_ids) if none) + """ + roles_tuple = tuple(target_roles) + cache_key = ( + id(tokenizer), + _messages_hash(messages), + roles_tuple, + content_pattern, + target_min_msg_idx, + ) + cached = _MASK_CACHE.get(cache_key) + if cached is not None: + return cached + + pattern = re.compile(content_pattern) if content_pattern else None + role_set = set(target_roles) + + input_ids, spans = _per_message_spans(tokenizer, messages) + mask = torch.zeros(len(input_ids), dtype=torch.bool) + first_target_idx = len(input_ids) + + for i, msg in enumerate(messages): + if i < target_min_msg_idx: + continue + if msg.get("role") not in role_set: + continue + if pattern is not None and not pattern.search(str(msg.get("content", ""))): + continue + start, end = spans[i] + if end > start: + mask[start:end] = True + if start < first_target_idx: + first_target_idx = start + + # Bound the cache. + if len(_MASK_CACHE) >= _MASK_CACHE_MAX: + # Drop an arbitrary entry (dict preserves insertion order; pop oldest). + _MASK_CACHE.pop(next(iter(_MASK_CACHE))) + _MASK_CACHE[cache_key] = (input_ids, mask, first_target_idx) + + return input_ids, mask, first_target_idx + + +def clear_mask_cache() -> None: + """Clear the target-mask cache (useful in tests).""" + _MASK_CACHE.clear() diff --git a/src/bakery/trainer.py b/src/bakery/trainer.py index d42e222..d143159 100644 --- a/src/bakery/trainer.py +++ b/src/bakery/trainer.py @@ -1,49 +1,55 @@ -"""Prompt baking trainer via KL divergence. +"""Context baking trainer via KL divergence. -Inherits from transformers.Trainer to get standard HF training infrastructure -(logging, checkpointing, schedulers, gradient accumulation) while implementing -the prompt baking training loop: +Generalizes the original prompt-baking approach to arbitrary prefix contexts — +system prompts, conversation histories, accumulated memories, few-shot examples. +The teacher sees the full prefix; the student sees a trimmed version (or none). +KL divergence distills the teacher into the student, baking the prefix into weights. -- Single model with PEFT adapter toggling (no duplicate weights) -- Teacher: adapters disabled, sees system prompt -- Student: adapters enabled, no system prompt -- Per-token masked KL divergence loss -- On-the-fly trajectory generation from teacher +Inherits from transformers.Trainer for standard HF infrastructure (logging, +checkpointing, schedulers, gradient accumulation). -Based on the Prompt Baking paper (arxiv 2409.13697). +Based on the Prompt Baking paper (arxiv 2409.13697), generalized. """ +from __future__ import annotations + import logging -import torch -from typing import Optional, Callable +import warnings +from typing import Callable, List, Optional +import torch from datasets import Dataset from transformers import ( - Trainer, + GenerationConfig, PreTrainedModel, PreTrainedTokenizerBase, + Trainer, TrainerCallback, - GenerationConfig, ) from transformers.trainer_utils import EvalPrediction -from bakery.kl import compute_kl_divergence, disable_adapters, padding_side +from bakery.kl import compute_kl_divergence, disable_adapters +from bakery.masking import build_target_mask logger = logging.getLogger(__name__) -class PromptBakingTrainer(Trainer): - """Trainer that bakes system prompts into model weights via KL divergence. +class ContextBakingTrainer(Trainer): + """Trainer that bakes an arbitrary prefix context into model weights via KL divergence. - Overrides compute_loss() and training_step() to inject prompt baking logic: - 1. training_step: generates trajectories from teacher (adapters disabled) - 2. compute_loss: computes per-token KL divergence between teacher and student + For each example: + - Teacher view: prefix_messages + turns (+ response) + - Student view: prefix_messages[-student_retained_turns:] + turns (+ response) + KL is computed on tokens belonging to messages whose role matches + `context_config.target_roles` and content matches `target_content_pattern` + (if set). Prefix tokens never receive loss. """ def __init__( self, model: PreTrainedModel | str | None = None, args=None, + context_config=None, train_dataset: Dataset | None = None, eval_dataset: Dataset | None = None, processing_class: PreTrainedTokenizerBase | None = None, @@ -55,12 +61,40 @@ def __init__( None, ), ): - self.system_prompt = args.system_prompt self.num_trajectories = args.num_trajectories self.trajectory_length = args.trajectory_length self.sampling_temperature = args.sampling_temperature self.kl_temperature = args.temperature + # Back-compat: if no ContextConfig is passed but args has system_prompt, + # auto-desugar into a minimal ContextConfig. Keeps direct-instantiation + # users on the old API working. + if context_config is None and getattr(args, "system_prompt", None): + from bakery.config import ContextConfig + + context_config = ContextConfig( + prefix_messages=[{"role": "system", "content": args.system_prompt}] + ) + + # Context configuration — prefix, student view, target mask. + self.context_config = context_config + self.prefix_messages: List[dict] = ( + list(context_config.prefix_messages) + if context_config and context_config.prefix_messages + else [] + ) + self.student_retained_turns: int = ( + context_config.student_retained_turns if context_config else 0 + ) + self.target_roles: List[str] = ( + list(context_config.target_roles) + if context_config and context_config.target_roles + else ["assistant"] + ) + self.target_content_pattern: Optional[str] = ( + context_config.target_content_pattern if context_config else None + ) + super().__init__( model=model, args=args, @@ -73,8 +107,6 @@ def __init__( optimizers=optimizers, ) - self._prompt_length_cache: dict[str, tuple[int, int]] = {} - self.model_accepts_loss_kwargs = False self.generation_config = GenerationConfig( max_new_tokens=self.trajectory_length, @@ -84,48 +116,270 @@ def __init__( pad_token_id=self.processing_class.pad_token_id, ) - # -- Chat formatting -- + # -- Tokenization -- - def _tokenize(self, text: str, **kwargs) -> dict: + def _tokenize(self, text, **kwargs) -> dict: max_len = getattr(self.args, "max_seq_length", None) if max_len: kwargs.setdefault("truncation", True) kwargs.setdefault("max_length", max_len) return self.processing_class(text, add_special_tokens=False, **kwargs) - def _get_prompt_lengths(self, user_message: str) -> tuple[int, int]: - """Return (teacher_prompt_length, student_prompt_length) for a user message. + # -- Message assembly -- + + def _row_prefix(self, row_prefix: Optional[List[dict]]) -> List[dict]: + """Return the effective prefix for a row: per-row if set, else global.""" + if row_prefix: + return list(row_prefix) + return list(self.prefix_messages) + + def _student_prefix(self, prefix: List[dict]) -> List[dict]: + """Trim the prefix down to the last `student_retained_turns` messages.""" + if self.student_retained_turns <= 0: + return [] + return prefix[-self.student_retained_turns :] - Results are cached because prompt lengths are deterministic for a given - user_message (the system_prompt is constant across the trainer's lifetime). + @staticmethod + def _append_response(messages: List[dict], response: Optional[str]) -> List[dict]: + if response: + return list(messages) + [{"role": "assistant", "content": response}] + return list(messages) + + def _build_example( + self, + row_prefix: Optional[List[dict]], + turns: List[dict], + response: Optional[str], + ): + """Build tokenized teacher + student views for a single example. + + Returns a dict with: + teacher_ids, teacher_mask (target mask over teacher tokens), + student_ids, student_mask (target mask over student tokens), + teacher_prefix_token_len, student_prefix_token_len + or None if no target tokens were found. """ - if user_message not in self._prompt_length_cache: - t_prompt = self._format_prompted(user_message) - t_len = self._tokenize(t_prompt, return_tensors="pt")["input_ids"].shape[1] - s_prompt = self._format_unprompted(user_message) - s_len = self._tokenize(s_prompt, return_tensors="pt")["input_ids"].shape[1] - self._prompt_length_cache[user_message] = (t_len, s_len) - return self._prompt_length_cache[user_message] - - def _format_prompted(self, user_message: str) -> str: - messages = [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": user_message}, - ] - return self.processing_class.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True + prefix = self._row_prefix(row_prefix) + student_prefix = self._student_prefix(prefix) + + teacher_messages = self._append_response(list(prefix) + list(turns), response) + student_messages = self._append_response( + list(student_prefix) + list(turns), response + ) + + # Target mask: prefix messages never count as targets. + teacher_ids, teacher_target_mask, teacher_first = build_target_mask( + self.processing_class, + teacher_messages, + self.target_roles, + self.target_content_pattern, + target_min_msg_idx=len(prefix), + ) + student_ids, student_target_mask, student_first = build_target_mask( + self.processing_class, + student_messages, + self.target_roles, + self.target_content_pattern, + target_min_msg_idx=len(student_prefix), ) - def _format_unprompted(self, user_message: str) -> str: - messages = [{"role": "user", "content": user_message}] - return self.processing_class.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True + if teacher_first >= len(teacher_ids) or student_first >= len(student_ids): + return None # no target tokens + + return { + "teacher_ids": teacher_ids, + "teacher_mask": teacher_target_mask, + "teacher_first": teacher_first, + "student_ids": student_ids, + "student_mask": student_target_mask, + "student_first": student_first, + } + + # -- Batch preparation -- + + @staticmethod + def _normalize_batch(inputs: dict) -> tuple[list, list, list]: + """Normalize batch into (prefix_messages_per_row, turns_per_row, responses). + + Accepts new-format batches (prefix_messages, turns, responses) and the + legacy format (user_messages, responses). Legacy format wraps each user + message as a single-turn [{role: user, content: msg}] list. + """ + if "turns" in inputs: + prefix_list = inputs.get("prefix_messages") or [None] * len(inputs["turns"]) + return ( + list(prefix_list), + [list(t) for t in inputs["turns"]], + list(inputs.get("responses", [None] * len(inputs["turns"]))), + ) + # Legacy shape + user_messages = inputs.get("user_messages", []) + responses = inputs.get("responses", [None] * len(user_messages)) + turns = [[{"role": "user", "content": m}] for m in user_messages] + prefix_list = [None] * len(user_messages) + return prefix_list, turns, list(responses) + + def _build_batch(self, inputs, model) -> Optional[dict]: + """Tokenize + pad the batch; return dict of tensors or None if empty. + + Returned dict: + teacher_fwd: kwargs for teacher forward (input_ids, attention_mask, ...) + student_fwd: kwargs for student forward + teacher_mask_padded: (B, T_t) bool mask of target tokens (aligned with padding) + student_mask_padded: (B, T_s) bool mask of target tokens + """ + prefix_list, turns_list, responses = self._normalize_batch(inputs) + if not turns_list: + return None + + built = [] + for prefix, turns, resp in zip(prefix_list, turns_list, responses): + # Require a response or a turn ending in assistant for KL to have a target. + if not resp and not any(m.get("role") == "assistant" for m in turns): + continue + # If response is empty string, skip. + if resp is not None and isinstance(resp, str) and not resp.strip(): + continue + example = self._build_example(prefix, turns, resp) + if example is None: + continue + built.append(example) + + if not built: + return None + + # Left-pad teacher and student id lists into tensors. + pad_id = self.processing_class.pad_token_id + if pad_id is None: + pad_id = self.processing_class.eos_token_id + + def _pad_left(seqs, masks): + max_len = max(len(s) for s in seqs) + out_ids = torch.full((len(seqs), max_len), pad_id, dtype=torch.long) + out_attn = torch.zeros((len(seqs), max_len), dtype=torch.long) + out_mask = torch.zeros((len(seqs), max_len), dtype=torch.bool) + for i, (ids, msk) in enumerate(zip(seqs, masks)): + n = len(ids) + out_ids[i, max_len - n :] = torch.tensor(ids, dtype=torch.long) + out_attn[i, max_len - n :] = 1 + out_mask[i, max_len - n :] = msk + return out_ids, out_attn, out_mask + + t_ids, t_attn, t_mask = _pad_left( + [b["teacher_ids"] for b in built], + [b["teacher_mask"] for b in built], + ) + s_ids, s_attn, s_mask = _pad_left( + [b["student_ids"] for b in built], + [b["student_mask"] for b in built], ) + device = model.device + t_ids = t_ids.to(device) + t_attn = t_attn.to(device) + s_ids = s_ids.to(device) + s_attn = s_attn.to(device) + # Keep masks on CPU until slicing; move to device per-sample in loss. + + teacher_fwd = {"input_ids": t_ids, "attention_mask": t_attn} + student_fwd = {"input_ids": s_ids, "attention_mask": s_attn} + self._inject_gemma_token_types(model, teacher_fwd, t_ids) + self._inject_gemma_token_types(model, student_fwd, s_ids) + + return { + "teacher_fwd": teacher_fwd, + "student_fwd": student_fwd, + "teacher_mask_padded": t_mask, + "student_mask_padded": s_mask, + "teacher_attn": t_attn, + "student_attn": s_attn, + } + + @staticmethod + def _inject_gemma_token_types(model, fwd: dict, input_ids: torch.Tensor) -> None: + """Gemma 3/4 require token_type_ids (and mm_token_type_ids for Gemma 4) + during training even for text-only inputs. All-zeros is correct.""" + mtype = getattr(model.config, "model_type", None) + if mtype in ("gemma3", "gemma4", "gemma4_text"): + fwd["token_type_ids"] = torch.zeros_like(input_ids) + if mtype in ("gemma4", "gemma4_text"): + fwd["mm_token_type_ids"] = torch.zeros_like(input_ids) + + # -- KL loss from aligned logits + masks -- + + def _kl_from_logits( + self, + teacher_logits: torch.Tensor, + student_logits: torch.Tensor, + teacher_mask_padded: torch.BoolTensor, + student_mask_padded: torch.BoolTensor, + ) -> Optional[torch.Tensor]: + """Compute mean per-example KL over aligned target tokens. + + Logits at position t predict token t+1, so we shift by 1. Teacher and + student may have different prefix lengths, so we align on the trainable + region by finding the first target token in each and taking the common + length. + """ + losses = [] + B = teacher_logits.shape[0] + device = student_logits.device + + for i in range(B): + t_mask = teacher_mask_padded[i].to(device) + s_mask = student_mask_padded[i].to(device) + + t_target_positions = t_mask.nonzero(as_tuple=False).squeeze(-1) + s_target_positions = s_mask.nonzero(as_tuple=False).squeeze(-1) + if t_target_positions.numel() == 0 or s_target_positions.numel() == 0: + continue + + t_start = int(t_target_positions[0].item()) + s_start = int(s_target_positions[0].item()) + + # Logits predict next token → shift by 1. + t_logits = teacher_logits[i : i + 1, t_start - 1 : -1, :] + s_logits = student_logits[i : i + 1, s_start - 1 : -1, :] + t_tail_mask = t_mask[t_start:].float().unsqueeze(0) + s_tail_mask = s_mask[s_start:].float().unsqueeze(0) + + min_len = min( + t_logits.shape[1], + s_logits.shape[1], + t_tail_mask.shape[1], + s_tail_mask.shape[1], + ) + if min_len == 0: + continue + + t_logits = t_logits[:, :min_len, :] + s_logits = s_logits[:, :min_len, :] + combined_mask = t_tail_mask[:, :min_len] * s_tail_mask[:, :min_len] + if combined_mask.sum() == 0: + continue + + loss = compute_kl_divergence( + t_logits.detach() if t_logits.requires_grad else t_logits, + s_logits, + combined_mask, + self.kl_temperature, + ) + losses.append(loss) + + if not losses: + return None + return torch.stack(losses).mean() + # -- Trajectory generation -- def _generate_trajectory(self, user_message: str) -> str: - prompt = self._format_prompted(user_message) + """Generate a response from the teacher (adapters disabled, full prefix visible).""" + teacher_messages = list(self.prefix_messages) + [ + {"role": "user", "content": user_message} + ] + prompt = self.processing_class.apply_chat_template( + teacher_messages, tokenize=False, add_generation_prompt=True + ) inputs = self._tokenize(prompt, return_tensors="pt").to(self.model.device) was_training = self.model.training @@ -145,274 +399,99 @@ def _generate_trajectory(self, user_message: str) -> str: ) return response.strip() - # -- Loss computation -- - - def _prepare_pairs(self, inputs): - """Extract and validate (user_message, response) pairs from inputs. - - Returns a list of (user_msg, response) tuples with empty responses - filtered out, or None if the batch is empty/invalid. - """ - user_messages = inputs.get("user_messages", []) - responses = inputs.get("responses", []) - if not user_messages or not responses: - return None - pairs = [ - (msg, resp) for msg, resp in zip(user_messages, responses) if resp.strip() - ] - return pairs if pairs else None - - def _build_texts_and_lengths(self, pairs): - """Build teacher/student chat texts and prompt lengths for each pair. - - Returns (teacher_texts, student_texts, teacher_prompt_lengths, - student_prompt_lengths). - """ - teacher_texts, student_texts = [], [] - teacher_prompt_lengths, student_prompt_lengths = [], [] - - for user_msg, response in pairs: - t_msgs = [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": user_msg}, - {"role": "assistant", "content": response}, - ] - teacher_texts.append( - self.processing_class.apply_chat_template(t_msgs, tokenize=False) - ) - - s_msgs = [ - {"role": "user", "content": user_msg}, - {"role": "assistant", "content": response}, - ] - student_texts.append( - self.processing_class.apply_chat_template(s_msgs, tokenize=False) - ) - - t_len, s_len = self._get_prompt_lengths(user_msg) - teacher_prompt_lengths.append(t_len) - student_prompt_lengths.append(s_len) - - return ( - teacher_texts, - student_texts, - teacher_prompt_lengths, - student_prompt_lengths, - ) + # -- Loss + eval -- - def _make_fwd_kwargs(self, model, tok_inputs): - """Build forward-pass keyword arguments, handling token_type_ids.""" - fwd = dict( - input_ids=tok_inputs["input_ids"], - attention_mask=tok_inputs["attention_mask"], - ) - if hasattr(model.config, "model_type") and model.config.model_type in ( - "gemma3", - ): - fwd["token_type_ids"] = torch.zeros_like(tok_inputs["input_ids"]) - elif "token_type_ids" in tok_inputs: - fwd["token_type_ids"] = tok_inputs["token_type_ids"] - return fwd - - def _compute_batched_kl( - self, - teacher_logits, - student_logits, - teacher_inputs, - student_inputs, - teacher_prompt_lengths, - student_prompt_lengths, - B, - ): - """Compute batched KL divergence from aligned teacher/student logits. - - Slices response-only logits from each sequence (accounting for - left-padding offsets), assembles them into aligned batch tensors, - and returns per-sample KL losses. - - Returns per-sample loss tensor of shape [|valid|], or None if no - valid aligned logit pairs exist. - """ - t_seq_len = teacher_inputs["input_ids"].shape[1] - s_seq_len = student_inputs["input_ids"].shape[1] - t_real_lengths = teacher_inputs["attention_mask"].sum(dim=1) - s_real_lengths = student_inputs["attention_mask"].sum(dim=1) - V = teacher_logits.shape[-1] - - # Compute per-sample response start positions (in logit space, shifted -1 - # so that logit[t] predicts token[t+1]). - t_starts = [ - int(t_seq_len - t_real_lengths[i].item()) + teacher_prompt_lengths[i] - for i in range(B) - ] - s_starts = [ - int(s_seq_len - s_real_lengths[i].item()) + student_prompt_lengths[i] - for i in range(B) - ] - # Response length for sample i: from start to seq_end (exclusive), capped - # at the other sequence's response length to keep teacher/student aligned. - t_resp_lens = [t_seq_len - t_starts[i] for i in range(B)] - s_resp_lens = [s_seq_len - s_starts[i] for i in range(B)] - min_resp_lens = [min(t_resp_lens[i], s_resp_lens[i]) for i in range(B)] - - # Filter out zero-length samples (degenerate prompts/responses). - valid = [i for i, L in enumerate(min_resp_lens) if L > 0] - if not valid: - return None - - max_resp_len = max(min_resp_lens[i] for i in valid) - - # Build batched logit tensors [|valid|, max_resp_len, V] by copying each - # sample's response slice. This CPU loop is cheap (shapes only differ in - # sequence position); the expensive softmax/KL runs once on the batch. - dev = student_logits.device - t_batch = student_logits.new_zeros(len(valid), max_resp_len, V) - s_batch = student_logits.new_zeros(len(valid), max_resp_len, V) - mask_batch = student_logits.new_zeros(len(valid), max_resp_len) - - for out_idx, i in enumerate(valid): - L = min_resp_lens[i] - ts = t_starts[i] - 1 # logit position for first response token - ss = s_starts[i] - 1 - t_batch[out_idx, :L] = teacher_logits[i, ts : ts + L].to(dev) - s_batch[out_idx, :L] = student_logits[i, ss : ss + L] - mask_batch[out_idx, :L] = 1.0 - - per_sample_losses = compute_kl_divergence( - t_batch.detach(), - s_batch, - mask_batch, - self.kl_temperature, - per_sample=True, - ) - return per_sample_losses + def _zero_loss(self) -> torch.Tensor: + return torch.tensor(0.0, device=self.args.device, requires_grad=True) def compute_loss( self, model, inputs, return_outputs=False, num_items_in_batch=None ): """Compute KL divergence loss with batched forward passes.""" - pairs = self._prepare_pairs(inputs) - if pairs is None: - logger.warning( - "Batch has no valid user_messages/responses — returning zero loss" - ) - loss = torch.tensor(0.0, device=self.args.device, requires_grad=True) - return (loss, None) if return_outputs else loss - - teacher_texts, student_texts, teacher_prompt_lengths, student_prompt_lengths = ( - self._build_texts_and_lengths(pairs) - ) - - with padding_side(self.processing_class, "left"): - teacher_inputs = self._tokenize( - teacher_texts, return_tensors="pt", padding=True - ).to(model.device) - student_inputs = self._tokenize( - student_texts, return_tensors="pt", padding=True - ).to(model.device) + batch = self._build_batch(inputs, model) + if batch is None: + logger.warning("Empty batch after building — returning zero loss") + zero = self._zero_loss() + return (zero, None) if return_outputs else zero with torch.no_grad(): with disable_adapters(model): - teacher_outputs = model(**self._make_fwd_kwargs(model, teacher_inputs)) - - student_outputs = model(**self._make_fwd_kwargs(model, student_inputs)) + teacher_outputs = model(**batch["teacher_fwd"]) + student_outputs = model(**batch["student_fwd"]) - per_sample_losses = self._compute_batched_kl( + loss = self._kl_from_logits( teacher_outputs.logits, student_outputs.logits, - teacher_inputs, - student_inputs, - teacher_prompt_lengths, - student_prompt_lengths, - len(pairs), + batch["teacher_mask_padded"], + batch["student_mask_padded"], ) - - if per_sample_losses is None: + if loss is None: logger.warning("No aligned logit pairs after slicing — returning zero loss") - zero = torch.tensor(0.0, device=self.args.device, requires_grad=True) + zero = self._zero_loss() return (zero, None) if return_outputs else zero - - total_loss = per_sample_losses.mean() - return (total_loss, None) if return_outputs else total_loss + return (loss, None) if return_outputs else loss def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): - """Eval step: reuse compute_loss so the collated batch format works. + """Eval step. With sequential_eval=True, teacher logits are moved to CPU before the student forward pass to halve peak VRAM usage. """ if not self.args.sequential_eval: - # Unwrap to bypass Accelerate's fp32 upcast wrapper raw = model.module if hasattr(model, "module") else model with torch.no_grad(): loss = self.compute_loss(raw, inputs) return (loss.detach(), None, None) - # Sequential eval: teacher → offload logits to CPU → student - pairs = self._prepare_pairs(inputs) - if not pairs: - return ( - torch.tensor(0.0, device=self.args.device, requires_grad=True), - None, - None, - ) - - teacher_texts, student_texts, teacher_prompt_lengths, student_prompt_lengths = ( - self._build_texts_and_lengths(pairs) - ) - - with padding_side(self.processing_class, "left"): - teacher_inputs = self._tokenize( - teacher_texts, return_tensors="pt", padding=True - ).to(model.device) - student_inputs = self._tokenize( - student_texts, return_tensors="pt", padding=True - ).to(model.device) + # Sequential: teacher → CPU offload → student + batch = self._build_batch(inputs, model) + if batch is None: + return (self._zero_loss(), None, None) - # Accelerate replaces model.forward with a wrapper that upcasts to fp32. - # Bypass by calling the CLASS forward method directly. base = model.module if hasattr(model, "module") else model fwd_fn = type(base).forward with torch.no_grad(): with disable_adapters(base): - teacher_logits = fwd_fn( - base, **self._make_fwd_kwargs(base, teacher_inputs) - ).logits.cpu() + teacher_logits = fwd_fn(base, **batch["teacher_fwd"]).logits.cpu() torch.cuda.empty_cache() - student_outputs = fwd_fn( - base, **self._make_fwd_kwargs(base, student_inputs) - ) + student_outputs = fwd_fn(base, **batch["student_fwd"]) - per_sample_losses = self._compute_batched_kl( - teacher_logits, + loss = self._kl_from_logits( + teacher_logits.to(student_outputs.logits.device), student_outputs.logits, - teacher_inputs, - student_inputs, - teacher_prompt_lengths, - student_prompt_lengths, - len(pairs), + batch["teacher_mask_padded"], + batch["student_mask_padded"], ) - - if per_sample_losses is None: - return ( - torch.tensor(0.0, device=self.args.device, requires_grad=True), - None, - None, - ) - - return (per_sample_losses.mean().detach(), None, None) + if loss is None: + return (self._zero_loss(), None, None) + return (loss.detach(), None, None) def training_step(self, model, inputs, num_items_in_batch=None) -> torch.Tensor: - """Generate trajectories on-the-fly if no precomputed responses.""" - existing_responses = inputs.get("responses", []) - if existing_responses: + """Generate trajectories on-the-fly if no precomputed responses. + + Trajectory mode requires the final turn to be `user` (we sample an + assistant response from the teacher). Multi-turn datasets with + explicit assistant turns skip this path and use the provided responses. + """ + existing_responses = inputs.get("responses") or [] + if any(r for r in existing_responses): return super().training_step(model, inputs, num_items_in_batch) - user_messages = inputs.get("user_messages", []) - all_user_messages, all_responses = [], [] + # Trajectory mode: need user_messages (legacy) or turns ending in user. + _, turns_list, _ = self._normalize_batch(inputs) + user_messages = [] + for turns in turns_list: + if not turns: + continue + last = turns[-1] + if last.get("role") != "user": + logger.warning( + "Trajectory mode: skipping row whose last turn is not 'user'" + ) + continue + user_messages.append(last["content"]) - # Threshold beyond which memory usage from accumulated trajectories - # may become significant (each trajectory requires a full forward pass). + all_user_messages, all_responses = [], [] _TRAJECTORY_WARN_THRESHOLD = 64 if len(user_messages) * self.num_trajectories > _TRAJECTORY_WARN_THRESHOLD: logger.warning( @@ -422,7 +501,6 @@ def training_step(self, model, inputs, num_items_in_batch=None) -> torch.Tensor: len(user_messages), self.num_trajectories, ) - for user_msg in user_messages: for _ in range(self.num_trajectories): response = self._generate_trajectory(user_msg) @@ -432,8 +510,28 @@ def training_step(self, model, inputs, num_items_in_batch=None) -> torch.Tensor: if not all_responses: logger.warning("No valid trajectories generated — returning zero loss") - return torch.tensor(0.0, device=self.args.device, requires_grad=True) + return self._zero_loss() + # Feed back as legacy-shape batch. inputs["user_messages"] = all_user_messages inputs["responses"] = all_responses + inputs.pop("turns", None) + inputs.pop("prefix_messages", None) return super().training_step(model, inputs, num_items_in_batch) + + +class PromptBakingTrainer(ContextBakingTrainer): + """Deprecated alias for ContextBakingTrainer. + + Bakery now generalizes prompt baking to arbitrary prefix contexts. + Use ContextBakingTrainer going forward. + """ + + def __init__(self, *args, **kwargs): + warnings.warn( + "PromptBakingTrainer is deprecated; use ContextBakingTrainer instead. " + "The old name will be removed in a future release.", + DeprecationWarning, + stacklevel=2, + ) + super().__init__(*args, **kwargs) diff --git a/tests/test_cli_helpers.py b/tests/test_cli_helpers.py new file mode 100644 index 0000000..cd75db1 --- /dev/null +++ b/tests/test_cli_helpers.py @@ -0,0 +1,95 @@ +"""Tests for CLI helper functions (_load_prefix_file).""" + +import json +import os +import tempfile + +import pytest + +from bakery.cli import _load_prefix_file + + +def _write(suffix: str, content: str) -> str: + f = tempfile.NamedTemporaryFile(mode="w", suffix=suffix, delete=False) + f.write(content) + f.close() + return f.name + + +def test_load_prefix_file_json(): + path = _write( + ".json", + json.dumps( + [ + {"role": "system", "content": "s"}, + {"role": "user", "content": "u"}, + ] + ), + ) + try: + result = _load_prefix_file(path) + assert result == [ + {"role": "system", "content": "s"}, + {"role": "user", "content": "u"}, + ] + finally: + os.unlink(path) + + +def test_load_prefix_file_yaml(): + path = _write( + ".yaml", + "- role: system\n content: s\n- role: user\n content: u\n", + ) + try: + result = _load_prefix_file(path) + assert len(result) == 2 + assert result[0]["role"] == "system" + assert result[1]["content"] == "u" + finally: + os.unlink(path) + + +def test_load_prefix_file_yml_extension(): + path = _write( + ".yml", + "- role: system\n content: hi\n", + ) + try: + result = _load_prefix_file(path) + assert result == [{"role": "system", "content": "hi"}] + finally: + os.unlink(path) + + +def test_load_prefix_file_rejects_non_list_dict(): + path = _write(".json", json.dumps({"role": "system", "content": "s"})) + try: + with pytest.raises(ValueError, match="must contain a JSON/YAML list"): + _load_prefix_file(path) + finally: + os.unlink(path) + + +def test_load_prefix_file_rejects_non_list_string(): + path = _write(".json", json.dumps("just a string")) + try: + with pytest.raises(ValueError): + _load_prefix_file(path) + finally: + os.unlink(path) + + +def test_load_prefix_file_empty_list(): + """An empty list is technically valid — caller may want to reject later.""" + path = _write(".json", json.dumps([])) + try: + result = _load_prefix_file(path) + assert result == [] + finally: + os.unlink(path) + + +def test_load_prefix_file_missing_file(): + with pytest.raises(FileNotFoundError): + _load_prefix_file("/nonexistent/path/prefix.json") diff --git a/tests/test_context_baking.py b/tests/test_context_baking.py new file mode 100644 index 0000000..ce47556 --- /dev/null +++ b/tests/test_context_baking.py @@ -0,0 +1,326 @@ +"""Tests for the context-baking generalization (prefix_messages, target masking, +multi-turn, per-row prefix override, retained student turns). +""" + +from peft import LoraConfig as PeftLoraConfig, get_peft_model +from transformers import AutoModelForCausalLM, AutoTokenizer + +from bakery.config import BakeryConfig, ContextConfig +from bakery.data import ( + create_conversational_dataset, + prompt_baking_collator, +) +from bakery.masking import build_target_mask, clear_mask_cache +from bakery.trainer import ContextBakingTrainer + +CHAT_TEMPLATE = ( + "{% for m in messages %}" + "{{ m['role'] }}: {{ m['content'] }}\n" + "{% endfor %}" + "{% if add_generation_prompt %}assistant: {% endif %}" +) + + +def _make_tokenizer(): + tok = AutoTokenizer.from_pretrained("gpt2") + tok.pad_token = tok.eos_token + tok.chat_template = CHAT_TEMPLATE + return tok + + +def _make_trainer( + context_config=None, + dataset=None, + system_prompt=None, + batch_size=1, +): + tokenizer = _make_tokenizer() + model = AutoModelForCausalLM.from_pretrained("gpt2") + peft_config = PeftLoraConfig( + r=4, + lora_alpha=8, + target_modules=["c_attn"], + task_type="CAUSAL_LM", + ) + model = get_peft_model(model, peft_config) + + args = BakeryConfig( + output_dir="/tmp/bakery_ctx_test", + system_prompt=system_prompt, + num_trajectories=1, + trajectory_length=8, + per_device_train_batch_size=batch_size, + num_train_epochs=1, + logging_steps=1, + report_to="none", + use_cpu=True, + ) + + return ContextBakingTrainer( + model=model, + args=args, + context_config=context_config, + train_dataset=dataset, + processing_class=tokenizer, + data_collator=prompt_baking_collator, + ) + + +# ---------- masking unit tests ---------- + + +def test_mask_role_based_single_assistant(): + """Single assistant turn: only its tokens are True.""" + clear_mask_cache() + tok = _make_tokenizer() + messages = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi there"}, + ] + ids, mask, first = build_target_mask(tok, messages, ["assistant"], None) + assert len(ids) == mask.numel() + assert mask.any() + assert 0 <= first < len(ids) + # Every True position must be in the assistant tail. + assert mask[:first].sum() == 0 + + +def test_mask_all_assistant_turns_multi_turn(): + """Multi-turn convo: every assistant message contributes to mask.""" + clear_mask_cache() + tok = _make_tokenizer() + messages = [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "a2"}, + ] + _, mask, _ = build_target_mask(tok, messages, ["assistant"], None) + # Should have two disjoint True regions (one per assistant). + diffs = mask.int().diff().abs().sum().item() + # Number of edges = 2 * num_spans. + assert diffs >= 2 + + +def test_mask_regex_filter(): + """Regex filters which assistant turns count as targets.""" + clear_mask_cache() + tok = _make_tokenizer() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "nope"}, + {"role": "user", "content": "q again"}, + {"role": "assistant", "content": "Answer: yes"}, + ] + _, mask_all, _ = build_target_mask(tok, messages, ["assistant"], None) + _, mask_filtered, _ = build_target_mask(tok, messages, ["assistant"], r"^Answer:") + assert mask_all.sum() > mask_filtered.sum() + assert mask_filtered.sum() > 0 + + +def test_mask_excludes_prefix_via_target_min_msg_idx(): + """target_min_msg_idx skips few-shot assistant turns in the prefix.""" + clear_mask_cache() + tok = _make_tokenizer() + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "fewshot-q"}, + {"role": "assistant", "content": "fewshot-a"}, # prefix, should be masked out + {"role": "user", "content": "real-q"}, + {"role": "assistant", "content": "real-a"}, + ] + _, mask_with_prefix, first_with = build_target_mask( + tok, messages, ["assistant"], None, target_min_msg_idx=3 + ) + _, mask_no_prefix, first_no = build_target_mask( + tok, messages, ["assistant"], None, target_min_msg_idx=0 + ) + assert mask_no_prefix.sum() > mask_with_prefix.sum() + assert first_with > first_no # first target is pushed later + + +def test_mask_cache_key_includes_prefix_idx(): + """Different target_min_msg_idx values produce distinct cache entries.""" + from bakery.masking import _MASK_CACHE + + clear_mask_cache() + tok = _make_tokenizer() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + build_target_mask(tok, messages, ["assistant"], None, target_min_msg_idx=0) + build_target_mask(tok, messages, ["assistant"], None, target_min_msg_idx=1) + assert len(_MASK_CACHE) == 2 + + +# ---------- trainer integration tests ---------- + + +def test_backcompat_system_prompt_desugars(): + """BakeryConfig.system_prompt with no ContextConfig still produces nonzero loss.""" + trainer = _make_trainer( + context_config=None, + system_prompt="You are a helpful assistant.", + ) + inputs = { + "user_messages": ["What is 2+2?"], + "responses": ["The answer is 4."], + } + loss = trainer.compute_loss(trainer.model, inputs) + # Desugared prefix means teacher sees system, student does not → real KL signal. + assert loss.item() > 0 + assert trainer.prefix_messages == [ + {"role": "system", "content": "You are a helpful assistant."} + ] + + +def test_global_prefix_messages_multi_turn(): + """Global multi-turn prefix (system + few-shot) produces a nonzero loss.""" + context = ContextConfig( + prefix_messages=[ + {"role": "system", "content": "You answer concisely."}, + {"role": "user", "content": "Example Q"}, + {"role": "assistant", "content": "Example A"}, + ], + target_roles=["assistant"], + ) + trainer = _make_trainer(context_config=context, system_prompt=None) + inputs = { + "user_messages": ["What is AI?"], + "responses": ["Artificial intelligence."], + } + loss = trainer.compute_loss(trainer.model, inputs) + assert loss.item() > 0 + + +def test_per_row_prefix_overrides_global(): + """A row's prefix_messages takes precedence over the global one.""" + global_prefix = [{"role": "system", "content": "global prompt"}] + row_prefix = [{"role": "system", "content": "row-specific prompt"}] + context = ContextConfig(prefix_messages=global_prefix) + trainer = _make_trainer(context_config=context) + + inputs = { + "prefix_messages": [row_prefix], + "turns": [[{"role": "user", "content": "hi"}]], + "responses": ["hello"], + } + loss = trainer.compute_loss(trainer.model, inputs) + # Sanity: this path uses the row prefix, not global. + assert loss.item() > 0 + + +def test_student_retained_turns_nonzero(): + """student_retained_turns=N keeps last N prefix messages in student view.""" + context_pure = ContextConfig( + prefix_messages=[ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "ctx-q"}, + {"role": "assistant", "content": "ctx-a"}, + ], + student_retained_turns=0, + ) + context_retained = ContextConfig( + prefix_messages=[ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "ctx-q"}, + {"role": "assistant", "content": "ctx-a"}, + ], + student_retained_turns=2, + ) + + t_pure = _make_trainer(context_config=context_pure) + t_retained = _make_trainer(context_config=context_retained) + + # Prefixes produce different student views → different prompt lengths. + assert t_pure.student_retained_turns == 0 + assert t_retained.student_retained_turns == 2 + assert t_pure._student_prefix(context_pure.prefix_messages) == [] + assert t_retained._student_prefix(context_retained.prefix_messages) == [ + {"role": "user", "content": "ctx-q"}, + {"role": "assistant", "content": "ctx-a"}, + ] + + +def test_multi_turn_conversational_batch(): + """Conversational batch with multi-turn turns → all assistant tokens are targets.""" + context = ContextConfig( + prefix_messages=[{"role": "system", "content": "sys"}], + ) + trainer = _make_trainer(context_config=context) + inputs = { + "prefix_messages": [None], + "turns": [ + [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "a2"}, + ] + ], + "responses": [None], + } + loss = trainer.compute_loss(trainer.model, inputs) + assert loss.item() > 0 + + +def test_target_roles_restrict_to_assistant(): + """Switching target_roles to exclude assistant yields zero-loss fallback.""" + context = ContextConfig( + prefix_messages=[{"role": "system", "content": "sys"}], + target_roles=["tool"], # no messages match + ) + trainer = _make_trainer(context_config=context) + inputs = { + "user_messages": ["q"], + "responses": ["a"], + } + loss = trainer.compute_loss(trainer.model, inputs) + assert loss.item() == 0.0 + + +def test_pattern_filter_selects_matching_turns_only(): + """target_content_pattern restricts targets to matching assistant content.""" + context = ContextConfig( + prefix_messages=[{"role": "system", "content": "sys"}], + target_roles=["assistant"], + target_content_pattern=r"^Answer:", + ) + trainer = _make_trainer(context_config=context) + # Response does not match pattern → zero loss. + inputs_nomatch = { + "user_messages": ["q"], + "responses": ["just talk"], + } + inputs_match = { + "user_messages": ["q"], + "responses": ["Answer: yes"], + } + assert trainer.compute_loss(trainer.model, inputs_nomatch).item() == 0.0 + assert trainer.compute_loss(trainer.model, inputs_match).item() > 0 + + +def test_create_conversational_dataset_shape(): + """create_conversational_dataset preserves prefix_messages/turns/response columns.""" + rows = [ + { + "turns": [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ], + "prefix_messages": [{"role": "system", "content": "s"}], + "response": None, + }, + { + "turns": [{"role": "user", "content": "q2"}], + "prefix_messages": None, + "response": "a2", + }, + ] + ds = create_conversational_dataset(rows) + assert set(ds.column_names) == {"turns", "prefix_messages", "responses"} + assert len(ds) == 2 + assert ds[0]["prefix_messages"] == [{"role": "system", "content": "s"}] + assert ds[1]["prefix_messages"] is None + assert ds[1]["responses"] == "a2" diff --git a/tests/test_context_config.py b/tests/test_context_config.py new file mode 100644 index 0000000..5b79f38 --- /dev/null +++ b/tests/test_context_config.py @@ -0,0 +1,91 @@ +"""Tests for ContextConfig dataclass and HfArgumentParser integration.""" + +from transformers import HfArgumentParser + +from bakery.config import BakeryConfig, ContextConfig, DataConfig, LoraConfig + + +def test_context_config_defaults(): + cfg = ContextConfig() + assert cfg.prefix_messages is None + assert cfg.prefix_messages_file is None + assert cfg.student_retained_turns == 0 + assert cfg.target_roles == ["assistant"] + assert cfg.target_content_pattern is None + + +def test_context_config_target_roles_not_shared_across_instances(): + """Mutable default: each instance should get its own list.""" + a = ContextConfig() + b = ContextConfig() + a.target_roles.append("user") + assert b.target_roles == ["assistant"] + + +def test_context_config_explicit_fields(): + cfg = ContextConfig( + prefix_messages=[{"role": "system", "content": "s"}], + student_retained_turns=2, + target_roles=["assistant", "tool"], + target_content_pattern=r"^A", + ) + assert cfg.prefix_messages == [{"role": "system", "content": "s"}] + assert cfg.student_retained_turns == 2 + assert cfg.target_roles == ["assistant", "tool"] + assert cfg.target_content_pattern == r"^A" + + +def test_hf_argument_parser_includes_context_fields(): + """ContextConfig fields are exposed as top-level CLI flags via HfArgumentParser.""" + parser = HfArgumentParser((BakeryConfig, DataConfig, LoraConfig, ContextConfig)) + baking, data, lora, context = parser.parse_args_into_dataclasses( + args=[ + "--output_dir", + "/tmp/hftest", + "--model_name_or_path", + "gpt2", + "--student_retained_turns", + "3", + "--target_content_pattern", + r"^Answer:", + ] + ) + assert context.student_retained_turns == 3 + assert context.target_content_pattern == r"^Answer:" + assert context.target_roles == ["assistant"] # default preserved + + +def test_bakery_config_system_prompt_still_accepted(): + """Deprecated system_prompt field still parseable for back-compat.""" + parser = HfArgumentParser((BakeryConfig, DataConfig, LoraConfig, ContextConfig)) + baking, _, _, _ = parser.parse_args_into_dataclasses( + args=[ + "--output_dir", + "/tmp/bp_test", + "--model_name_or_path", + "gpt2", + "--system_prompt", + "You are helpful.", + ] + ) + assert baking.system_prompt == "You are helpful." + + +def test_bakery_config_and_context_config_coexist(): + """Both old system_prompt and new prefix_messages fields can be set simultaneously + — CLI layer is responsible for reconciling / warning.""" + parser = HfArgumentParser((BakeryConfig, DataConfig, LoraConfig, ContextConfig)) + baking, _, _, context = parser.parse_args_into_dataclasses( + args=[ + "--output_dir", + "/tmp/both", + "--model_name_or_path", + "gpt2", + "--system_prompt", + "old-style prompt", + "--student_retained_turns", + "1", + ] + ) + assert baking.system_prompt == "old-style prompt" + assert context.student_retained_turns == 1 diff --git a/tests/test_data_conversational.py b/tests/test_data_conversational.py new file mode 100644 index 0000000..8d52e10 --- /dev/null +++ b/tests/test_data_conversational.py @@ -0,0 +1,256 @@ +"""Tests for conversational data loading and the generalized collator.""" + +import json +import os +import tempfile + +import pytest + +from bakery.data import ( + create_conversational_dataset, + load_conversations, + prompt_baking_collator, +) + + +# ---------- create_conversational_dataset ---------- + + +def test_create_conversational_dataset_empty(): + ds = create_conversational_dataset([]) + assert len(ds) == 0 + assert set(ds.column_names) == {"turns", "prefix_messages", "responses"} + + +def test_create_conversational_dataset_single_turn_with_response(): + ds = create_conversational_dataset( + [ + { + "turns": [{"role": "user", "content": "hi"}], + "prefix_messages": None, + "response": "hello", + } + ] + ) + assert len(ds) == 1 + assert ds[0]["turns"] == [{"role": "user", "content": "hi"}] + assert ds[0]["prefix_messages"] is None + assert ds[0]["responses"] == "hello" + + +def test_create_conversational_dataset_missing_keys_default_to_none_or_empty(): + """Rows that omit optional fields still produce a well-formed dataset.""" + ds = create_conversational_dataset([{"turns": []}]) + assert len(ds) == 1 + assert ds[0]["turns"] == [] + assert ds[0]["prefix_messages"] is None + assert ds[0]["responses"] is None + + +def test_create_conversational_dataset_preserves_row_order(): + rows = [ + {"turns": [{"role": "user", "content": f"q{i}"}], "response": f"a{i}"} + for i in range(5) + ] + ds = create_conversational_dataset(rows) + for i in range(5): + assert ds[i]["turns"][0]["content"] == f"q{i}" + assert ds[i]["responses"] == f"a{i}" + + +def test_create_conversational_dataset_mixed_prefix_presence(): + rows = [ + { + "turns": [{"role": "user", "content": "a"}], + "prefix_messages": [{"role": "system", "content": "P"}], + }, + {"turns": [{"role": "user", "content": "b"}]}, + ] + ds = create_conversational_dataset(rows) + assert ds[0]["prefix_messages"] == [{"role": "system", "content": "P"}] + assert ds[1]["prefix_messages"] is None + + +# ---------- prompt_baking_collator ---------- + + +def test_collator_legacy_shape(): + feats = [ + {"user_messages": "q1", "responses": "a1"}, + {"user_messages": "q2", "responses": "a2"}, + ] + batch = prompt_baking_collator(feats) + assert "turns" not in batch + assert batch["user_messages"] == ["q1", "q2"] + assert batch["responses"] == ["a1", "a2"] + + +def test_collator_conversational_shape(): + feats = [ + { + "turns": [{"role": "user", "content": "q"}], + "prefix_messages": [{"role": "system", "content": "s"}], + "responses": "a", + } + ] + batch = prompt_baking_collator(feats) + assert batch["turns"] == [[{"role": "user", "content": "q"}]] + assert batch["prefix_messages"] == [[{"role": "system", "content": "s"}]] + assert batch["responses"] == ["a"] + assert "user_messages" not in batch + + +def test_collator_preserves_none_prefix(): + feats = [ + { + "turns": [{"role": "user", "content": "q"}], + "prefix_messages": None, + "responses": None, + } + ] + batch = prompt_baking_collator(feats) + assert batch["prefix_messages"] == [None] + assert batch["responses"] == [None] + + +def test_collator_conversational_handles_null_turns(): + """None turns should be coerced to [].""" + feats = [{"turns": None, "prefix_messages": None, "responses": None}] + batch = prompt_baking_collator(feats) + assert batch["turns"] == [[]] + + +def test_collator_batch_size_multi_conversational(): + feats = [ + { + "turns": [{"role": "user", "content": f"q{i}"}], + "prefix_messages": None, + "responses": f"a{i}", + } + for i in range(3) + ] + batch = prompt_baking_collator(feats) + assert len(batch["turns"]) == 3 + assert len(batch["responses"]) == 3 + + +def test_collator_empty_legacy_features(): + batch = prompt_baking_collator([]) + # Empty legacy shape is the default when no features have 'turns'. + assert batch == {"user_messages": [], "responses": []} + + +def test_collator_legacy_list_values_are_extended(): + feats = [{"user_messages": ["a", "b"], "responses": ["x", "y"]}] + batch = prompt_baking_collator(feats) + assert batch["user_messages"] == ["a", "b"] + assert batch["responses"] == ["x", "y"] + + +# ---------- load_conversations (JSON) ---------- + + +@pytest.fixture +def tmpjson(): + paths = [] + + def _write(obj): + f = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) + json.dump(obj, f) + f.close() + paths.append(f.name) + return f.name + + yield _write + + for p in paths: + try: + os.unlink(p) + except OSError: + pass + + +def test_load_conversations_json_messages_format(tmpjson): + path = tmpjson( + [ + { + "messages": [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + ] + } + ] + ) + rows = load_conversations(path) + assert len(rows) == 1 + assert rows[0]["turns"] == [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + ] + assert rows[0]["prefix_messages"] is None + assert rows[0]["response"] is None + + +def test_load_conversations_json_with_per_row_prefix(tmpjson): + path = tmpjson( + [ + { + "prefix_messages": [{"role": "system", "content": "persona A"}], + "messages": [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ], + } + ] + ) + rows = load_conversations(path) + assert rows[0]["prefix_messages"] == [{"role": "system", "content": "persona A"}] + + +def test_load_conversations_json_prompt_response_pairs(tmpjson): + path = tmpjson( + [ + {"prompt": "q1", "response": "a1"}, + {"prompt": "q2", "response": "a2"}, + ] + ) + rows = load_conversations(path) + assert len(rows) == 2 + assert rows[0]["turns"] == [{"role": "user", "content": "q1"}] + assert rows[0]["response"] == "a1" + + +def test_load_conversations_json_plain_string_list(tmpjson): + """List of strings becomes single-turn user-only rows.""" + path = tmpjson(["q1", "q2"]) + rows = load_conversations(path) + assert len(rows) == 2 + assert rows[0]["turns"] == [{"role": "user", "content": "q1"}] + assert rows[0]["response"] is None + assert rows[0]["prefix_messages"] is None + + +def test_load_conversations_json_nested_data_key(tmpjson): + path = tmpjson({"conversations": [{"prompt": "q", "response": "a"}]}) + rows = load_conversations(path) + assert len(rows) == 1 + assert rows[0]["turns"] == [{"role": "user", "content": "q"}] + assert rows[0]["response"] == "a" + + +def test_load_conversations_json_skips_empty_prompts(tmpjson): + """Rows with empty prompt + no messages column should be skipped.""" + path = tmpjson([{"prompt": "", "response": "a"}, {"prompt": "q", "response": "a"}]) + rows = load_conversations(path) + # At least the valid row is present. + assert any(r["turns"] == [{"role": "user", "content": "q"}] for r in rows) + + +def test_load_conversations_json_alternate_prompt_keys(tmpjson): + path = tmpjson([{"question": "q", "answer": "a"}]) + rows = load_conversations(path) + assert rows[0]["turns"] == [{"role": "user", "content": "q"}] + # `answer` isn't in the recognized response keys, so response is None. + # (Only prompt-side key remapping is aggressive; response uses canonical keys.) + # We just assert the turn was loaded correctly. + assert len(rows) == 1 diff --git a/tests/test_format_trajectory.py b/tests/test_format_trajectory.py deleted file mode 100644 index 3986159..0000000 --- a/tests/test_format_trajectory.py +++ /dev/null @@ -1,327 +0,0 @@ -"""Tests for PromptBakingTrainer format helpers and _generate_trajectory.""" - -import pytest -from unittest.mock import patch, MagicMock -from transformers import AutoModelForCausalLM, AutoTokenizer -from peft import LoraConfig as PeftLoraConfig, get_peft_model - -from bakery.config import BakeryConfig -from bakery.data import create_dataset, prompt_baking_collator -from bakery.trainer import PromptBakingTrainer - - -CHAT_TEMPLATE = ( - "{% for m in messages %}" - "{{ m['role'] }}: {{ m['content'] }}\n" - "{% endfor %}" - "{% if add_generation_prompt %}assistant: {% endif %}" -) - - -def _make_trainer(system_prompt="Be helpful.", num_trajectories=1): - tokenizer = AutoTokenizer.from_pretrained("gpt2") - tokenizer.pad_token = tokenizer.eos_token - tokenizer.chat_template = CHAT_TEMPLATE - - model = AutoModelForCausalLM.from_pretrained("gpt2") - peft_config = PeftLoraConfig( - r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" - ) - model = get_peft_model(model, peft_config) - - args = BakeryConfig( - output_dir="/tmp/bakery_test", - system_prompt=system_prompt, - num_trajectories=num_trajectories, - trajectory_length=16, - per_device_train_batch_size=1, - num_train_epochs=1, - logging_steps=1, - report_to="none", - use_cpu=True, - ) - dataset = create_dataset(["Hello?"]) - trainer = PromptBakingTrainer( - model=model, - args=args, - train_dataset=dataset, - processing_class=tokenizer, - data_collator=prompt_baking_collator, - ) - return trainer - - -# --------------------------------------------------------------------------- -# _format_prompted -# --------------------------------------------------------------------------- - - -class TestFormatPrompted: - def test_includes_system_prompt(self): - trainer = _make_trainer(system_prompt="Custom system.") - result = trainer._format_prompted("Hello?") - assert "Custom system." in result - - def test_includes_user_message(self): - trainer = _make_trainer() - result = trainer._format_prompted("What is the capital of France?") - assert "What is the capital of France?" in result - - def test_includes_generation_prompt(self): - """apply_chat_template with add_generation_prompt=True appends 'assistant: '.""" - trainer = _make_trainer() - result = trainer._format_prompted("Hello?") - assert "assistant:" in result - - def test_returns_string(self): - trainer = _make_trainer() - assert isinstance(trainer._format_prompted("Hi"), str) - - -# --------------------------------------------------------------------------- -# _format_unprompted -# --------------------------------------------------------------------------- - - -class TestFormatUnprompted: - def test_excludes_system_prompt(self): - trainer = _make_trainer(system_prompt="System content here.") - result = trainer._format_unprompted("Hello?") - assert "System content here." not in result - - def test_includes_user_message(self): - trainer = _make_trainer() - result = trainer._format_unprompted("What time is it?") - assert "What time is it?" in result - - def test_includes_generation_prompt(self): - trainer = _make_trainer() - result = trainer._format_unprompted("Hello?") - assert "assistant:" in result - - def test_shorter_than_prompted(self): - """Unprompted text should be shorter than prompted (no system msg).""" - trainer = _make_trainer(system_prompt="A fairly long system prompt.") - msg = "Hello?" - assert len(trainer._format_unprompted(msg)) < len(trainer._format_prompted(msg)) - - -# --------------------------------------------------------------------------- -# _generate_trajectory -# --------------------------------------------------------------------------- - - -class TestGenerateTrajectory: - def test_returns_string(self): - trainer = _make_trainer() - result = trainer._generate_trajectory("What is 2+2?") - assert isinstance(result, str) - - def test_strips_whitespace(self): - """_generate_trajectory must strip leading/trailing whitespace.""" - trainer = _make_trainer() - # Patch decode to return padded string - with patch.object( - trainer.processing_class, - "decode", - return_value=" Some response. ", - ): - result = trainer._generate_trajectory("Q?") - assert result == "Some response." - - def test_model_switched_to_eval_during_generation(self): - """Model is set to eval() during generation.""" - trainer = _make_trainer() - trainer.model.train() - eval_modes = [] - - original_generate = trainer.model.generate - - def spy_generate(**kwargs): - eval_modes.append(trainer.model.training) - return original_generate(**kwargs) - - with patch.object(trainer.model, "generate", side_effect=spy_generate): - trainer._generate_trajectory("Q?") - - assert eval_modes and eval_modes[0] is False - - def test_model_restored_to_train_after_generation(self): - """Model is restored to train() after generation.""" - trainer = _make_trainer() - trainer.model.train() - trainer._generate_trajectory("Q?") - assert trainer.model.training - - def test_model_stays_eval_if_was_eval(self): - """If model was already in eval(), it stays in eval() after generation.""" - trainer = _make_trainer() - trainer.model.eval() - trainer._generate_trajectory("Q?") - assert not trainer.model.training - - def test_adapters_disabled_during_generation(self): - """_generate_trajectory uses disable_adapters context manager.""" - trainer = _make_trainer() - with patch("bakery.trainer.disable_adapters") as mock_ctx: - mock_cm = MagicMock() - mock_cm.__enter__ = MagicMock(return_value=None) - mock_cm.__exit__ = MagicMock(return_value=False) - mock_ctx.return_value = mock_cm - trainer._generate_trajectory("Q?") - mock_ctx.assert_called_once() - - -# --------------------------------------------------------------------------- -# _tokenize -# --------------------------------------------------------------------------- - - -class TestTokenize: - def test_tokenize_returns_dict_with_input_ids(self): - trainer = _make_trainer() - result = trainer._tokenize("Hello world", return_tensors="pt") - assert "input_ids" in result - - def test_tokenize_no_special_tokens(self): - """add_special_tokens=False is always set.""" - trainer = _make_trainer() - # With special tokens the beginning-of-sequence token would be prepended. - result_no_special = trainer._tokenize("Hello", return_tensors="pt") - # Directly call the tokenizer with add_special_tokens=True for comparison - with_special = trainer.processing_class( - "Hello", add_special_tokens=True, return_tensors="pt" - ) - # Without special tokens should produce ≤ tokens than with - assert ( - result_no_special["input_ids"].shape[1] - <= with_special["input_ids"].shape[1] - ) - - def test_tokenize_respects_max_seq_length(self): - """When max_seq_length is set, output is truncated.""" - trainer = _make_trainer() - trainer.args.max_seq_length = 3 - long_text = "one two three four five six seven eight nine ten" - result = trainer._tokenize(long_text, return_tensors="pt") - assert result["input_ids"].shape[1] <= 3 - - def test_tokenize_no_truncation_without_max_seq_length(self): - """When max_seq_length is None, long text is not truncated.""" - trainer = _make_trainer() - trainer.args.max_seq_length = None - long_text = " ".join(["word"] * 50) - result = trainer._tokenize(long_text, return_tensors="pt") - # Should be more than 10 tokens - assert result["input_ids"].shape[1] > 10 - - def test_tokenize_accepts_list_of_strings(self): - """Batch tokenization with a list works.""" - trainer = _make_trainer() - result = trainer._tokenize( - ["Hello", "World"], return_tensors="pt", padding=True - ) - assert result["input_ids"].shape[0] == 2 - - -# --------------------------------------------------------------------------- -# kl_temperature is taken from args.temperature -# --------------------------------------------------------------------------- - - -class TestKLTemperature: - def test_kl_temperature_matches_args_temperature(self): - """trainer.kl_temperature equals args.temperature at construction.""" - trainer = _make_trainer() - assert trainer.kl_temperature == trainer.args.temperature - - def test_kl_temperature_custom_value(self): - """Custom temperature is reflected in kl_temperature.""" - tokenizer = AutoTokenizer.from_pretrained("gpt2") - tokenizer.pad_token = tokenizer.eos_token - tokenizer.chat_template = CHAT_TEMPLATE - - model = AutoModelForCausalLM.from_pretrained("gpt2") - peft_config = PeftLoraConfig( - r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" - ) - model = get_peft_model(model, peft_config) - - args = BakeryConfig( - output_dir="/tmp/bakery_test", - system_prompt="You are a helper.", - num_trajectories=1, - trajectory_length=16, - per_device_train_batch_size=1, - num_train_epochs=1, - logging_steps=1, - report_to="none", - use_cpu=True, - temperature=2.5, - ) - from bakery.data import create_dataset, prompt_baking_collator - - dataset = create_dataset(["Q?"]) - trainer = PromptBakingTrainer( - model=model, - args=args, - train_dataset=dataset, - processing_class=tokenizer, - data_collator=prompt_baking_collator, - ) - assert trainer.kl_temperature == pytest.approx(2.5) - - -# --------------------------------------------------------------------------- -# generation_config -# --------------------------------------------------------------------------- - - -class TestGenerationConfig: - def test_generation_config_max_new_tokens(self): - trainer = _make_trainer() - assert trainer.generation_config.max_new_tokens == trainer.trajectory_length - - def test_generation_config_temperature(self): - trainer = _make_trainer() - assert trainer.generation_config.temperature == pytest.approx( - trainer.sampling_temperature - ) - - def test_generation_config_do_sample(self): - trainer = _make_trainer() - assert trainer.generation_config.do_sample is True - - def test_generation_config_pad_token_id(self): - trainer = _make_trainer() - assert ( - trainer.generation_config.pad_token_id - == trainer.processing_class.pad_token_id - ) - - -# --------------------------------------------------------------------------- -# __init__ attributes -# --------------------------------------------------------------------------- - - -class TestInitAttributes: - def test_system_prompt_stored(self): - trainer = _make_trainer(system_prompt="Custom prompt.") - assert trainer.system_prompt == "Custom prompt." - - def test_num_trajectories_stored(self): - trainer = _make_trainer(num_trajectories=3) - assert trainer.num_trajectories == 3 - - def test_trajectory_length_stored(self): - trainer = _make_trainer() - assert trainer.trajectory_length == 16 # from _make_trainer fixture - - def test_sampling_temperature_stored(self): - trainer = _make_trainer() - assert trainer.sampling_temperature == trainer.args.sampling_temperature - - def test_prompt_length_cache_starts_empty(self): - trainer = _make_trainer() - assert trainer._prompt_length_cache == {} diff --git a/tests/test_masking.py b/tests/test_masking.py new file mode 100644 index 0000000..4033d00 --- /dev/null +++ b/tests/test_masking.py @@ -0,0 +1,388 @@ +"""Comprehensive tests for bakery.masking.build_target_mask. + +Covers: +- Role filtering (single + multiple roles) +- Regex filtering +- target_min_msg_idx (prefix exclusion) +- Cache behavior (hits, misses, key components, eviction, clearing) +- Edge cases: empty messages, no matches, all matches +- Span ordering / non-overlap sanity +""" + +import pytest +import torch +from transformers import AutoTokenizer + +from bakery.masking import ( + _MASK_CACHE, + _MASK_CACHE_MAX, + _longest_common_prefix_len, + _messages_hash, + _per_message_spans, + build_target_mask, + clear_mask_cache, +) + +CHAT_TEMPLATE = ( + "{% for m in messages %}" + "{{ m['role'] }}: {{ m['content'] }}\n" + "{% endfor %}" + "{% if add_generation_prompt %}assistant: {% endif %}" +) + + +@pytest.fixture +def tok(): + tokenizer = AutoTokenizer.from_pretrained("gpt2") + tokenizer.pad_token = tokenizer.eos_token + tokenizer.chat_template = CHAT_TEMPLATE + return tokenizer + + +# ---------- helpers ---------- + + +def test_longest_common_prefix_len_identical(): + assert _longest_common_prefix_len([1, 2, 3], [1, 2, 3]) == 3 + + +def test_longest_common_prefix_len_partial(): + assert _longest_common_prefix_len([1, 2, 3, 4], [1, 2, 9]) == 2 + + +def test_longest_common_prefix_len_empty(): + assert _longest_common_prefix_len([], [1, 2]) == 0 + assert _longest_common_prefix_len([1, 2], []) == 0 + assert _longest_common_prefix_len([], []) == 0 + + +def test_longest_common_prefix_len_no_overlap(): + assert _longest_common_prefix_len([9, 8], [1, 2]) == 0 + + +def test_messages_hash_stable(): + m = [{"role": "user", "content": "x"}] + assert _messages_hash(m) == _messages_hash(m) + + +def test_messages_hash_changes_with_content(): + a = [{"role": "user", "content": "x"}] + b = [{"role": "user", "content": "y"}] + assert _messages_hash(a) != _messages_hash(b) + + +def test_messages_hash_changes_with_role(): + a = [{"role": "user", "content": "x"}] + b = [{"role": "assistant", "content": "x"}] + assert _messages_hash(a) != _messages_hash(b) + + +def test_messages_hash_handles_missing_fields(): + # Malformed messages shouldn't crash the hash. + _messages_hash([{"role": "user"}]) + _messages_hash([{"content": "x"}]) + _messages_hash([{}]) + + +def test_per_message_spans_covers_full_sequence(tok): + """Sum of distinct span ranges should cover every token in the full sequence.""" + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + ] + full_ids, spans = _per_message_spans(tok, messages) + # Spans should be non-decreasing in start and end. + for i in range(len(spans) - 1): + assert spans[i][1] <= spans[i + 1][1] + # Final span should reach full length. + assert spans[-1][1] == len(full_ids) + # Starts should never exceed ends. + for s, e in spans: + assert s <= e + + +def test_per_message_spans_empty(tok): + full_ids, spans = _per_message_spans(tok, []) + assert full_ids == [] + assert spans == [] + + +# ---------- build_target_mask ---------- + + +def test_mask_returns_list_tensor_and_index(tok): + clear_mask_cache() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "there"}, + ] + ids, mask, first = build_target_mask(tok, messages, ["assistant"]) + assert isinstance(ids, list) + assert isinstance(mask, torch.Tensor) + assert mask.dtype == torch.bool + assert mask.numel() == len(ids) + assert isinstance(first, int) + + +def test_mask_no_target_role_match(tok): + """When no message matches target_roles, mask is all False and first == len.""" + clear_mask_cache() + messages = [ + {"role": "user", "content": "only user"}, + {"role": "system", "content": "only system"}, + ] + ids, mask, first = build_target_mask(tok, messages, ["assistant"]) + assert mask.sum().item() == 0 + assert first == len(ids) + + +def test_mask_target_role_user(tok): + """target_roles=['user'] should mask user tokens, not assistant.""" + clear_mask_cache() + messages = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + _, mask_user, _ = build_target_mask(tok, messages, ["user"]) + _, mask_assistant, _ = build_target_mask(tok, messages, ["assistant"]) + assert mask_user.sum() > 0 + assert mask_assistant.sum() > 0 + # Roles are mutually exclusive — no shared True positions. + assert (mask_user & mask_assistant).sum() == 0 + + +def test_mask_multiple_target_roles(tok): + """Multiple roles in target_roles → union of their tokens.""" + clear_mask_cache() + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "u"}, + {"role": "assistant", "content": "a"}, + ] + _, mask_user, _ = build_target_mask(tok, messages, ["user"]) + _, mask_asst, _ = build_target_mask(tok, messages, ["assistant"]) + _, mask_both, _ = build_target_mask(tok, messages, ["user", "assistant"]) + assert mask_both.sum() == mask_user.sum() + mask_asst.sum() + + +def test_mask_regex_matches_all(tok): + """Pattern matching every content → same as no pattern.""" + clear_mask_cache() + messages = [ + {"role": "assistant", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + _, m_all, _ = build_target_mask(tok, messages, ["assistant"], None) + _, m_re, _ = build_target_mask(tok, messages, ["assistant"], r".") + assert m_all.sum() == m_re.sum() + + +def test_mask_regex_matches_none(tok): + """Pattern matching nothing → zero mask.""" + clear_mask_cache() + messages = [ + {"role": "assistant", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + _, mask, first = build_target_mask( + tok, messages, ["assistant"], r"zzzzzzzz_definitely_not" + ) + assert mask.sum() == 0 + assert first == len(mask) + + +def test_mask_regex_re_search_semantics(tok): + """Pattern uses re.search (not re.fullmatch) — partial match should count.""" + clear_mask_cache() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "prefix Answer: yes suffix"}, + ] + _, mask, _ = build_target_mask(tok, messages, ["assistant"], r"Answer:") + assert mask.sum() > 0 + + +def test_mask_first_target_idx_points_to_true_position(tok): + clear_mask_cache() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + _, mask, first = build_target_mask(tok, messages, ["assistant"]) + assert first < len(mask) + assert mask[first].item() is True + + +def test_mask_first_target_idx_empty_when_no_targets(tok): + clear_mask_cache() + messages = [{"role": "user", "content": "q"}] + ids, mask, first = build_target_mask(tok, messages, ["assistant"]) + assert first == len(ids) + + +def test_mask_target_min_msg_idx_past_end_gives_empty(tok): + """target_min_msg_idx >= len(messages) → no targets possible.""" + clear_mask_cache() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + _, mask, first = build_target_mask( + tok, messages, ["assistant"], None, target_min_msg_idx=5 + ) + assert mask.sum() == 0 + assert first == len(mask) + + +def test_mask_target_min_msg_idx_zero_equals_default(tok): + clear_mask_cache() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + _, m_default, _ = build_target_mask(tok, messages, ["assistant"]) + _, m_zero, _ = build_target_mask( + tok, messages, ["assistant"], None, target_min_msg_idx=0 + ) + assert m_default.sum() == m_zero.sum() + + +def test_mask_prefix_assistant_excluded_as_target(tok): + """Few-shot assistant in prefix must NOT be a target when target_min_msg_idx is set.""" + clear_mask_cache() + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "ex-q"}, + {"role": "assistant", "content": "ex-a"}, # few-shot prefix answer + {"role": "user", "content": "real-q"}, + {"role": "assistant", "content": "real-a"}, # training target + ] + ids, mask, first = build_target_mask( + tok, messages, ["assistant"], None, target_min_msg_idx=3 + ) + # Only one contiguous True region. + transitions = mask.int().diff().abs().sum().item() + assert transitions <= 2 # 0→1 and 1→0 + assert mask.sum() > 0 + + +# ---------- cache ---------- + + +def test_cache_hit_on_repeat(tok): + clear_mask_cache() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + build_target_mask(tok, messages, ["assistant"]) + assert len(_MASK_CACHE) == 1 + build_target_mask(tok, messages, ["assistant"]) + # Still exactly one entry. + assert len(_MASK_CACHE) == 1 + + +def test_cache_key_differentiates_target_roles(tok): + clear_mask_cache() + messages = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}] + build_target_mask(tok, messages, ["assistant"]) + build_target_mask(tok, messages, ["user"]) + build_target_mask(tok, messages, ["assistant", "user"]) + assert len(_MASK_CACHE) == 3 + + +def test_cache_key_differentiates_pattern(tok): + clear_mask_cache() + messages = [{"role": "assistant", "content": "Answer: 42"}] + build_target_mask(tok, messages, ["assistant"], None) + build_target_mask(tok, messages, ["assistant"], r"^Answer:") + build_target_mask(tok, messages, ["assistant"], r"42$") + assert len(_MASK_CACHE) == 3 + + +def test_cache_key_differentiates_messages(tok): + clear_mask_cache() + build_target_mask( + tok, + [{"role": "user", "content": "q1"}, {"role": "assistant", "content": "a"}], + ["assistant"], + ) + build_target_mask( + tok, + [{"role": "user", "content": "q2"}, {"role": "assistant", "content": "a"}], + ["assistant"], + ) + assert len(_MASK_CACHE) == 2 + + +def test_cache_eviction_bounded(tok, monkeypatch): + """When cache is at max, inserting a new entry evicts the oldest.""" + clear_mask_cache() + monkeypatch.setattr("bakery.masking._MASK_CACHE_MAX", 3) + for i in range(10): + msgs = [ + {"role": "user", "content": f"q{i}"}, + {"role": "assistant", "content": f"a{i}"}, + ] + build_target_mask(tok, msgs, ["assistant"]) + # Cache size should never exceed the bound (module-level _MASK_CACHE_MAX + # may have been re-read but the check uses the patched value). + assert len(_MASK_CACHE) <= 3 + + +def test_clear_mask_cache(tok): + messages = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}] + build_target_mask(tok, messages, ["assistant"]) + assert len(_MASK_CACHE) > 0 + clear_mask_cache() + assert len(_MASK_CACHE) == 0 + + +def test_mask_cache_max_constant_exists(): + """Ensure the bound constant is a sensible default.""" + assert isinstance(_MASK_CACHE_MAX, int) + assert _MASK_CACHE_MAX >= 100 + + +# ---------- structural invariants ---------- + + +def test_mask_always_bool_tensor(tok): + clear_mask_cache() + for messages in [ + [{"role": "user", "content": "x"}], + [{"role": "assistant", "content": "x"}], + [ + {"role": "system", "content": "s"}, + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ], + ]: + _, mask, _ = build_target_mask(tok, messages, ["assistant"]) + assert mask.dtype == torch.bool + + +def test_mask_length_matches_input_ids(tok): + clear_mask_cache() + messages = [ + {"role": "user", "content": "longer query here"}, + {"role": "assistant", "content": "a longer response with more tokens in it"}, + ] + ids, mask, _ = build_target_mask(tok, messages, ["assistant"]) + assert mask.numel() == len(ids) + + +def test_mask_target_region_is_contiguous_for_single_assistant(tok): + """A single assistant message produces exactly one contiguous True region.""" + clear_mask_cache() + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "some answer"}, + ] + _, mask, _ = build_target_mask(tok, messages, ["assistant"]) + if mask.sum() == 0: + pytest.skip("no targets — tokenizer template anomaly") + transitions = mask.int().diff().abs().sum().item() + # Starts at False, rises to True once, falls back at most once. + assert transitions in (1, 2) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 5d9d517..5e92bb1 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1,12 +1,10 @@ """Tests for PromptBakingTrainer using a tiny GPT-2 model with LoRA.""" -import torch from transformers import AutoModelForCausalLM, AutoTokenizer from peft import LoraConfig as PeftLoraConfig, get_peft_model from bakery.config import BakeryConfig from bakery.data import create_dataset, prompt_baking_collator -from bakery.kl import compute_kl_divergence, disable_adapters, padding_side from bakery.trainer import PromptBakingTrainer @@ -154,31 +152,36 @@ def test_loss_is_differentiable(): assert has_grad, "No LoRA parameters received gradients" -def test_prompt_length_cache(): - """Prompt lengths are cached and reused across compute_loss calls.""" +def test_target_mask_cache_grows_per_unique_conversation(): + """Target-mask cache accumulates one entry per unique (teacher, student) convo.""" + from bakery.masking import _MASK_CACHE, clear_mask_cache + + clear_mask_cache() trainer = _make_trainer( prompts=["What is 2+2?"], responses=["The answer is 4."], ) - assert len(trainer._prompt_length_cache) == 0 + assert len(_MASK_CACHE) == 0 inputs = { "user_messages": ["What is 2+2?"], "responses": ["The answer is 4."], } trainer.compute_loss(trainer.model, inputs) - assert "What is 2+2?" in trainer._prompt_length_cache - t_len, s_len = trainer._prompt_length_cache["What is 2+2?"] - assert isinstance(t_len, int) and t_len > 0 - assert isinstance(s_len, int) and s_len > 0 + # Two entries: teacher-view and student-view of the single conversation. + size_after_first = len(_MASK_CACHE) + assert size_after_first >= 1 - # Second call with same prompt should reuse cache (no new entries) + # Second call with identical inputs should not add new entries. trainer.compute_loss(trainer.model, inputs) - assert len(trainer._prompt_length_cache) == 1 + assert len(_MASK_CACHE) == size_after_first + +def test_target_mask_cache_multiple_prompts(): + """Different prompts produce distinct cache entries.""" + from bakery.masking import _MASK_CACHE, clear_mask_cache -def test_prompt_length_cache_multiple_prompts(): - """Cache accumulates entries for different prompts.""" + clear_mask_cache() trainer = _make_trainer( prompts=["Q1", "Q2"], responses=["A1", "A2"], @@ -189,9 +192,8 @@ def test_prompt_length_cache_multiple_prompts(): "responses": ["A1", "A2"], } trainer.compute_loss(trainer.model, inputs) - assert len(trainer._prompt_length_cache) == 2 - assert "Q1" in trainer._prompt_length_cache - assert "Q2" in trainer._prompt_length_cache + # At least 2 unique conversations → at least 2 entries. + assert len(_MASK_CACHE) >= 2 # --------------------------------------------------------------------------- @@ -245,7 +247,7 @@ def test_prediction_step_sequential_eval_returns_triple(): from transformers import AutoModelForCausalLM, AutoTokenizer from peft import LoraConfig as PeftLoraConfig, get_peft_model from bakery.data import create_dataset, prompt_baking_collator - from bakery.trainer import PromptBakingTrainer + from bakery.trainer import ContextBakingTrainer tokenizer = AutoTokenizer.from_pretrained("gpt2") tokenizer.pad_token = tokenizer.eos_token @@ -271,7 +273,7 @@ def test_prediction_step_sequential_eval_returns_triple(): use_cpu=True, sequential_eval=True, ) - trainer = PromptBakingTrainer( + trainer = ContextBakingTrainer( model=model, args=args, train_dataset=create_dataset(["Q"], ["A"]), @@ -285,100 +287,3 @@ def test_prediction_step_sequential_eval_returns_triple(): assert result[1] is None and result[2] is None loss = result[0] assert loss.dim() == 0 - - -# --------------------------------------------------------------------------- -# Numerical equivalence: batched vs per-sample loop -# --------------------------------------------------------------------------- - - -def test_batched_kl_matches_per_sample_loop(): - """Verify that the batched _compute_batched_kl produces the same result - as computing KL divergence one sample at a time in a loop. - - This guards against regressions when refactoring the vectorized path. - """ - torch.manual_seed(42) - - trainer = _make_trainer( - prompts=["What is 2+2?", "Explain gravity"], - responses=["The answer is 4.", "Gravity is a fundamental force of nature."], - batch_size=2, - ) - model = trainer.model - - user_messages = ["What is 2+2?", "Explain gravity"] - responses = ["The answer is 4.", "Gravity is a fundamental force of nature."] - pairs = list(zip(user_messages, responses)) - - teacher_texts, student_texts, t_prompt_lens, s_prompt_lens = ( - trainer._build_texts_and_lengths(pairs) - ) - - with padding_side(trainer.processing_class, "left"): - teacher_inputs = trainer._tokenize( - teacher_texts, return_tensors="pt", padding=True - ).to(model.device) - student_inputs = trainer._tokenize( - student_texts, return_tensors="pt", padding=True - ).to(model.device) - - with torch.no_grad(): - with disable_adapters(model): - teacher_logits = model( - **trainer._make_fwd_kwargs(model, teacher_inputs) - ).logits - student_logits = model(**trainer._make_fwd_kwargs(model, student_inputs)).logits - - # --- Batched path (the code under test) --- - batched_losses = trainer._compute_batched_kl( - teacher_logits, - student_logits, - teacher_inputs, - student_inputs, - t_prompt_lens, - s_prompt_lens, - len(pairs), - ) - assert batched_losses is not None - - # --- Reference: per-sample loop (the old approach) --- - per_sample_losses = [] - t_seq_len = teacher_inputs["input_ids"].shape[1] - s_seq_len = student_inputs["input_ids"].shape[1] - t_real_lengths = teacher_inputs["attention_mask"].sum(dim=1) - s_real_lengths = student_inputs["attention_mask"].sum(dim=1) - - for i in range(len(pairs)): - t_start = int(t_seq_len - t_real_lengths[i].item()) + t_prompt_lens[i] - s_start = int(s_seq_len - s_real_lengths[i].item()) + s_prompt_lens[i] - t_resp_len = t_seq_len - t_start - s_resp_len = s_seq_len - s_start - L = min(t_resp_len, s_resp_len) - if L <= 0: - continue - - ts = t_start - 1 # logit position for first response token - ss = s_start - 1 - t_logits_i = teacher_logits[i, ts : ts + L].unsqueeze(0) - s_logits_i = student_logits[i, ss : ss + L].unsqueeze(0) - mask_i = torch.ones(1, L, device=model.device) - - loss_i = compute_kl_divergence( - t_logits_i.detach(), - s_logits_i, - mask_i, - trainer.kl_temperature, - per_sample=True, - ) - per_sample_losses.append(loss_i.squeeze(0)) - - assert len(per_sample_losses) == batched_losses.shape[0] - loop_losses = torch.stack(per_sample_losses) - - assert torch.allclose(batched_losses, loop_losses, atol=1e-5), ( - f"Batched and per-sample loop KL losses differ:\n" - f" batched: {batched_losses}\n" - f" loop: {loop_losses}\n" - f" max diff: {(batched_losses - loop_losses).abs().max().item()}" - ) diff --git a/tests/test_trainer_internals.py b/tests/test_trainer_internals.py new file mode 100644 index 0000000..a989a35 --- /dev/null +++ b/tests/test_trainer_internals.py @@ -0,0 +1,510 @@ +"""Tests for ContextBakingTrainer internal helpers + edge-case integration tests. + +Uses one shared tiny GPT-2 + LoRA trainer across many tests to keep CPU time down. +""" + +import warnings + +import pytest +from peft import LoraConfig as PeftLoraConfig, get_peft_model +from transformers import AutoModelForCausalLM, AutoTokenizer + +from bakery.config import BakeryConfig, ContextConfig +from bakery.data import prompt_baking_collator +from bakery.masking import clear_mask_cache +from bakery.trainer import ContextBakingTrainer, PromptBakingTrainer + +CHAT_TEMPLATE = ( + "{% for m in messages %}" + "{{ m['role'] }}: {{ m['content'] }}\n" + "{% endfor %}" + "{% if add_generation_prompt %}assistant: {% endif %}" +) + + +def _mk_tokenizer(): + tok = AutoTokenizer.from_pretrained("gpt2") + tok.pad_token = tok.eos_token + tok.chat_template = CHAT_TEMPLATE + return tok + + +def _mk_trainer(context_config=None, system_prompt=None, batch_size=1): + tokenizer = _mk_tokenizer() + model = AutoModelForCausalLM.from_pretrained("gpt2") + peft_config = PeftLoraConfig( + r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" + ) + model = get_peft_model(model, peft_config) + args = BakeryConfig( + output_dir="/tmp/bakery_internals", + system_prompt=system_prompt, + num_trajectories=1, + trajectory_length=8, + per_device_train_batch_size=batch_size, + num_train_epochs=1, + logging_steps=1, + report_to="none", + use_cpu=True, + ) + return ContextBakingTrainer( + model=model, + args=args, + context_config=context_config, + processing_class=tokenizer, + data_collator=prompt_baking_collator, + ) + + +@pytest.fixture(scope="module") +def trainer_sys(): + """Trainer with a simple system prefix, reused across many tests.""" + return _mk_trainer( + context_config=ContextConfig( + prefix_messages=[{"role": "system", "content": "You are helpful."}] + ) + ) + + +@pytest.fixture(scope="module") +def trainer_multi_prefix(): + """Trainer with a multi-turn prefix (system + few-shot).""" + return _mk_trainer( + context_config=ContextConfig( + prefix_messages=[ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "ex-q"}, + {"role": "assistant", "content": "ex-a"}, + ], + student_retained_turns=0, + ) + ) + + +# ---------- _row_prefix ---------- + + +def test_row_prefix_uses_global_when_row_is_none(trainer_sys): + result = trainer_sys._row_prefix(None) + assert result == trainer_sys.prefix_messages + assert result is not trainer_sys.prefix_messages # defensive copy + + +def test_row_prefix_uses_global_when_row_is_empty(trainer_sys): + result = trainer_sys._row_prefix([]) + assert result == trainer_sys.prefix_messages + + +def test_row_prefix_prefers_row_over_global(trainer_sys): + row = [{"role": "system", "content": "row-specific"}] + result = trainer_sys._row_prefix(row) + assert result == row + + +# ---------- _student_prefix ---------- + + +def test_student_prefix_zero_returns_empty(): + t = _mk_trainer( + context_config=ContextConfig( + prefix_messages=[{"role": "system", "content": "s"}], + student_retained_turns=0, + ) + ) + assert t._student_prefix(t.prefix_messages) == [] + + +def test_student_prefix_n_larger_than_prefix_returns_full(): + prefix = [ + {"role": "system", "content": "s"}, + {"role": "user", "content": "u"}, + ] + t = _mk_trainer( + context_config=ContextConfig( + prefix_messages=prefix, + student_retained_turns=10, + ) + ) + assert t._student_prefix(prefix) == prefix + + +def test_student_prefix_keeps_last_n(): + prefix = [ + {"role": "system", "content": "s"}, + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + ] + t = _mk_trainer( + context_config=ContextConfig( + prefix_messages=prefix, + student_retained_turns=2, + ) + ) + assert t._student_prefix(prefix) == prefix[-2:] + + +# ---------- _append_response ---------- + + +def test_append_response_with_string(): + out = ContextBakingTrainer._append_response([{"role": "user", "content": "q"}], "a") + assert out[-1] == {"role": "assistant", "content": "a"} + assert len(out) == 2 + + +def test_append_response_with_none_passthrough(): + msgs = [{"role": "user", "content": "q"}] + out = ContextBakingTrainer._append_response(msgs, None) + assert out == msgs + assert out is not msgs # defensive copy + + +def test_append_response_with_empty_string_passthrough(): + """Empty/falsy response should not be appended.""" + msgs = [{"role": "user", "content": "q"}] + out = ContextBakingTrainer._append_response(msgs, "") + assert out == msgs + + +# ---------- _normalize_batch ---------- + + +def test_normalize_batch_legacy_shape(): + inputs = {"user_messages": ["q1", "q2"], "responses": ["a1", "a2"]} + prefix, turns, responses = ContextBakingTrainer._normalize_batch(inputs) + assert prefix == [None, None] + assert turns == [ + [{"role": "user", "content": "q1"}], + [{"role": "user", "content": "q2"}], + ] + assert responses == ["a1", "a2"] + + +def test_normalize_batch_conversational_shape(): + inputs = { + "turns": [[{"role": "user", "content": "q"}]], + "prefix_messages": [[{"role": "system", "content": "s"}]], + "responses": ["a"], + } + prefix, turns, responses = ContextBakingTrainer._normalize_batch(inputs) + assert prefix == [[{"role": "system", "content": "s"}]] + assert turns == [[{"role": "user", "content": "q"}]] + assert responses == ["a"] + + +def test_normalize_batch_conversational_missing_prefix_defaults_to_none(): + """When `turns` is present but `prefix_messages` column is missing / falsy, fill with None.""" + inputs = {"turns": [[{"role": "user", "content": "q"}]], "responses": ["a"]} + prefix, turns, responses = ContextBakingTrainer._normalize_batch(inputs) + assert prefix == [None] + + +def test_normalize_batch_legacy_empty(): + prefix, turns, responses = ContextBakingTrainer._normalize_batch( + {"user_messages": [], "responses": []} + ) + assert prefix == [] + assert turns == [] + assert responses == [] + + +# ---------- _build_example ---------- + + +def test_build_example_returns_none_when_no_target(trainer_sys): + """Prefix-only messages (no turns, no response) → no target tokens.""" + clear_mask_cache() + result = trainer_sys._build_example( + row_prefix=None, + turns=[], + response=None, + ) + assert result is None + + +def test_build_example_basic_single_turn(trainer_sys): + clear_mask_cache() + result = trainer_sys._build_example( + row_prefix=None, + turns=[{"role": "user", "content": "q"}], + response="a", + ) + assert result is not None + assert "teacher_ids" in result + assert "student_ids" in result + assert result["teacher_mask"].sum() > 0 + assert result["student_mask"].sum() > 0 + + +def test_build_example_prefix_assistant_not_a_target(trainer_multi_prefix): + """Few-shot assistant in the prefix must not become a training target.""" + clear_mask_cache() + result = trainer_multi_prefix._build_example( + row_prefix=None, + turns=[{"role": "user", "content": "real-q"}], + response="real-a", + ) + assert result is not None + # Exactly one target region (just the response) — the few-shot assistant is baked. + transitions = result["teacher_mask"].int().diff().abs().sum().item() + assert transitions <= 2 + + +# ---------- compute_loss edge cases ---------- + + +def test_compute_loss_no_response_and_no_assistant_turn_returns_zero(trainer_sys): + """A user-only turn with no response or assistant in the turns must not crash.""" + loss = trainer_sys.compute_loss( + trainer_sys.model, + { + "prefix_messages": [None], + "turns": [[{"role": "user", "content": "q"}]], + "responses": [None], + }, + ) + assert loss.item() == 0.0 + + +def test_compute_loss_whitespace_response_returns_zero(trainer_sys): + loss = trainer_sys.compute_loss( + trainer_sys.model, + {"user_messages": ["q"], "responses": [" \n\t"]}, + ) + assert loss.item() == 0.0 + + +def test_compute_loss_mixed_batch_some_valid(trainer_sys): + """Batch where one row is invalid (empty response) and one is valid → nonzero loss.""" + loss = trainer_sys.compute_loss( + trainer_sys.model, + {"user_messages": ["q1", "q2"], "responses": [" ", "a valid response"]}, + ) + assert loss.item() > 0 + + +def test_compute_loss_return_outputs_tuple(trainer_sys): + out = trainer_sys.compute_loss( + trainer_sys.model, + {"user_messages": ["q"], "responses": ["a"]}, + return_outputs=True, + ) + assert isinstance(out, tuple) + assert len(out) == 2 + assert out[1] is None + + +def test_compute_loss_return_outputs_zero_case(trainer_sys): + """return_outputs branch with no targets still returns a tuple (not scalar).""" + out = trainer_sys.compute_loss( + trainer_sys.model, + {"user_messages": [], "responses": []}, + return_outputs=True, + ) + assert isinstance(out, tuple) + assert out[0].item() == 0.0 + assert out[1] is None + + +def test_compute_loss_is_differentiable(trainer_sys): + loss = trainer_sys.compute_loss( + trainer_sys.model, {"user_messages": ["q"], "responses": ["a response"]} + ) + loss.backward() + # At least one LoRA param got a gradient. + has_grad = any( + p.grad is not None and p.requires_grad for p in trainer_sys.model.parameters() + ) + assert has_grad + + +def test_compute_loss_conversational_multi_turn_all_assistant_are_targets(trainer_sys): + """Multi-turn conversation with two assistant messages → nonzero loss over both spans.""" + loss = trainer_sys.compute_loss( + trainer_sys.model, + { + "prefix_messages": [None], + "turns": [ + [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "a2"}, + ] + ], + "responses": [None], + }, + ) + assert loss.item() > 0 + + +def test_compute_loss_target_roles_user_does_not_crash(): + """Reconfiguring target_roles to include 'user' must run without errors. + + Note: depending on student_retained_turns, the teacher/student alignment may + not overlap on user-role tokens (student lacks the prefix), so loss can be 0. + We only assert the path runs and returns a non-negative scalar. + """ + t = _mk_trainer( + context_config=ContextConfig( + prefix_messages=[{"role": "system", "content": "s"}], + target_roles=["user"], + ) + ) + loss = t.compute_loss( + t.model, + { + "prefix_messages": [None], + "turns": [ + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "there"}, + ] + ], + "responses": [None], + }, + ) + assert loss.item() >= 0 + + +def test_compute_loss_per_row_prefix_overrides_global(trainer_sys): + """A per-row prefix takes precedence over the trainer's global prefix.""" + row_prefix = [{"role": "system", "content": "row-specific"}] + loss = trainer_sys.compute_loss( + trainer_sys.model, + { + "prefix_messages": [row_prefix], + "turns": [[{"role": "user", "content": "q"}]], + "responses": ["a"], + }, + ) + assert loss.item() > 0 + + +def test_compute_loss_with_student_retained_turns_nonzero(): + """student_retained_turns > 0 still yields a real KL signal.""" + t = _mk_trainer( + context_config=ContextConfig( + prefix_messages=[ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "ctx-q"}, + {"role": "assistant", "content": "ctx-a"}, + ], + student_retained_turns=2, + ) + ) + loss = t.compute_loss(t.model, {"user_messages": ["hi"], "responses": ["there"]}) + # With retained turns student already sees ctx — loss is smaller but non-negative. + assert loss.item() >= 0 + + +# ---------- back-compat ---------- + + +def test_prompt_baking_trainer_is_deprecated_alias(): + tokenizer = _mk_tokenizer() + model = AutoModelForCausalLM.from_pretrained("gpt2") + peft_config = PeftLoraConfig( + r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" + ) + model = get_peft_model(model, peft_config) + args = BakeryConfig( + output_dir="/tmp/bakery_deprecation", + system_prompt="legacy prompt", + per_device_train_batch_size=1, + num_train_epochs=1, + logging_steps=1, + report_to="none", + use_cpu=True, + ) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + t = PromptBakingTrainer( + model=model, + args=args, + processing_class=tokenizer, + data_collator=prompt_baking_collator, + ) + assert isinstance(t, ContextBakingTrainer) + assert any( + issubclass(item.category, DeprecationWarning) + and "PromptBakingTrainer" in str(item.message) + for item in w + ) + + +def test_back_compat_system_prompt_auto_desugars_in_trainer_init(): + """args.system_prompt set but no context_config → trainer auto-wraps it.""" + tokenizer = _mk_tokenizer() + model = AutoModelForCausalLM.from_pretrained("gpt2") + peft_config = PeftLoraConfig( + r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" + ) + model = get_peft_model(model, peft_config) + args = BakeryConfig( + output_dir="/tmp/bakery_desugar", + system_prompt="Auto desugared!", + per_device_train_batch_size=1, + num_train_epochs=1, + logging_steps=1, + report_to="none", + use_cpu=True, + ) + t = ContextBakingTrainer( + model=model, + args=args, + processing_class=tokenizer, + data_collator=prompt_baking_collator, + ) + assert t.prefix_messages == [{"role": "system", "content": "Auto desugared!"}] + + +# ---------- prediction_step ---------- + + +def test_prediction_step_returns_scalar_loss(trainer_sys): + """Non-sequential eval path returns a detached scalar.""" + loss, _, _ = trainer_sys.prediction_step( + trainer_sys.model, + {"user_messages": ["q"], "responses": ["a"]}, + prediction_loss_only=True, + ) + assert loss.dim() == 0 + # Detached — no grad tracking. + assert not loss.requires_grad + + +def test_prediction_step_empty_inputs_returns_zero(trainer_sys): + loss, _, _ = trainer_sys.prediction_step( + trainer_sys.model, + {"user_messages": [], "responses": []}, + prediction_loss_only=True, + ) + assert loss.item() == 0.0 + + +# ---------- module exports ---------- + + +def test_module_exports_top_level_symbols(): + import bakery + + expected = { + "BakeryConfig", + "ContextConfig", + "DataConfig", + "LoraConfig", + "ContextBakingTrainer", + "PromptBakingTrainer", + "create_conversational_dataset", + "create_dataset", + "load_conversations", + "load_dataset", + "prompt_baking_collator", + "compute_kl_divergence", + "build_target_mask", + } + assert expected.issubset(set(bakery.__all__)) + for name in expected: + assert getattr(bakery, name) is not None diff --git a/tests/test_training_step.py b/tests/test_training_step.py deleted file mode 100644 index b9d7aa7..0000000 --- a/tests/test_training_step.py +++ /dev/null @@ -1,220 +0,0 @@ -"""Tests for PromptBakingTrainer.training_step.""" - -import logging - -import torch -from unittest.mock import patch -from transformers import AutoModelForCausalLM, AutoTokenizer -from peft import LoraConfig as PeftLoraConfig, get_peft_model - -from bakery.config import BakeryConfig -from bakery.data import create_dataset, prompt_baking_collator -from bakery.trainer import PromptBakingTrainer - - -CHAT_TEMPLATE = ( - "{% for m in messages %}" - "{{ m['role'] }}: {{ m['content'] }}\n" - "{% endfor %}" - "{% if add_generation_prompt %}assistant: {% endif %}" -) - - -def _make_trainer(prompts=None, responses=None, num_trajectories=1, batch_size=1): - """Tiny GPT-2 + LoRA trainer for testing.""" - prompts = prompts or ["What is 2+2?"] - tokenizer = AutoTokenizer.from_pretrained("gpt2") - tokenizer.pad_token = tokenizer.eos_token - tokenizer.chat_template = CHAT_TEMPLATE - - model = AutoModelForCausalLM.from_pretrained("gpt2") - peft_config = PeftLoraConfig( - r=4, - lora_alpha=8, - target_modules=["c_attn"], - task_type="CAUSAL_LM", - ) - model = get_peft_model(model, peft_config) - - args = BakeryConfig( - output_dir="/tmp/bakery_test", - system_prompt="You are a helpful assistant.", - num_trajectories=num_trajectories, - trajectory_length=16, - per_device_train_batch_size=batch_size, - num_train_epochs=1, - logging_steps=1, - report_to="none", - use_cpu=True, - ) - dataset = create_dataset(prompts, responses) - trainer = PromptBakingTrainer( - model=model, - args=args, - train_dataset=dataset, - processing_class=tokenizer, - data_collator=prompt_baking_collator, - ) - return trainer - - -# --------------------------------------------------------------------------- -# training_step: responses already present → delegates to super -# --------------------------------------------------------------------------- - - -def test_training_step_with_existing_responses_delegates_to_super(): - """When inputs already have responses, training_step calls super().""" - trainer = _make_trainer() - inputs = { - "user_messages": ["What is 2+2?"], - "responses": ["The answer is 4."], - } - with patch.object( - PromptBakingTrainer.__bases__[0], - "training_step", - return_value=torch.tensor(1.5), - ) as mock_super: - loss = trainer.training_step(trainer.model, inputs) - mock_super.assert_called_once() - assert loss.item() == 1.5 - - -# --------------------------------------------------------------------------- -# training_step: no responses → generates trajectories via _generate_trajectory -# --------------------------------------------------------------------------- - - -def test_training_step_generates_trajectories_when_no_responses(): - """When no responses present, training_step calls _generate_trajectory.""" - trainer = _make_trainer(num_trajectories=2) - inputs = {"user_messages": ["Hello?"]} - - with ( - patch.object( - trainer, "_generate_trajectory", return_value="Hi there!" - ) as mock_gen, - patch.object( - PromptBakingTrainer.__bases__[0], - "training_step", - return_value=torch.tensor(0.5), - ), - ): - trainer.training_step(trainer.model, inputs) - - # Called once per prompt × num_trajectories - assert mock_gen.call_count == 2 - - -def test_training_step_populates_inputs_with_generated_responses(): - """training_step fills inputs['responses'] before delegating to super.""" - trainer = _make_trainer(num_trajectories=1) - inputs = {"user_messages": ["What colour is the sky?"]} - - captured_inputs = {} - - def fake_super(model, inputs_arg, num_items=None): - captured_inputs.update(inputs_arg) - return torch.tensor(0.0) - - with ( - patch.object(trainer, "_generate_trajectory", return_value="Blue."), - patch.object( - PromptBakingTrainer.__bases__[0], - "training_step", - side_effect=fake_super, - ), - ): - trainer.training_step(trainer.model, inputs) - - assert "responses" in captured_inputs - assert captured_inputs["responses"] == ["Blue."] - - -def test_training_step_filters_blank_trajectories(): - """Blank trajectories are filtered out before delegating to super.""" - trainer = _make_trainer(num_trajectories=3) - inputs = {"user_messages": ["Question?"]} - - # Return blank for first two, valid for third - side_effects = ["", " ", "Valid answer."] - - with ( - patch.object(trainer, "_generate_trajectory", side_effect=side_effects), - patch.object( - PromptBakingTrainer.__bases__[0], - "training_step", - return_value=torch.tensor(0.0), - ) as mock_super, - ): - trainer.training_step(trainer.model, inputs) - - call_inputs = mock_super.call_args[0][1] - assert len(call_inputs["responses"]) == 1 - assert call_inputs["responses"][0] == "Valid answer." - - -def test_training_step_returns_zero_when_all_trajectories_blank(): - """All-blank trajectories → zero-loss tensor, super not called.""" - trainer = _make_trainer(num_trajectories=2) - inputs = {"user_messages": ["Prompt"]} - - with ( - patch.object(trainer, "_generate_trajectory", return_value=""), - patch.object( - PromptBakingTrainer.__bases__[0], - "training_step", - ) as mock_super, - ): - loss = trainer.training_step(trainer.model, inputs) - - mock_super.assert_not_called() - assert loss.item() == 0.0 - assert loss.requires_grad - - -# --------------------------------------------------------------------------- -# training_step: warning when many trajectories -# --------------------------------------------------------------------------- - - -def test_training_step_warns_when_many_trajectories(caplog): - """Logs a warning when total trajectories exceed threshold (64).""" - # 9 prompts × 8 trajectories = 72 > 64 threshold - prompts = [f"p{i}" for i in range(9)] - trainer = _make_trainer(prompts=prompts, num_trajectories=8) - inputs = {"user_messages": prompts} - - with ( - patch.object(trainer, "_generate_trajectory", return_value="ok"), - patch.object( - PromptBakingTrainer.__bases__[0], - "training_step", - return_value=torch.tensor(0.0), - ), - caplog.at_level(logging.WARNING), - ): - trainer.training_step(trainer.model, inputs) - - assert any("trajectories" in r.message.lower() for r in caplog.records) - - -def test_training_step_no_warn_below_threshold(caplog): - """Does NOT warn when total trajectories are at or below 64.""" - # 4 prompts × 16 trajectories = 64, exactly at threshold, no warning - prompts = [f"p{i}" for i in range(4)] - trainer = _make_trainer(prompts=prompts, num_trajectories=16) - inputs = {"user_messages": prompts} - - with ( - patch.object(trainer, "_generate_trajectory", return_value="ok"), - patch.object( - PromptBakingTrainer.__bases__[0], - "training_step", - return_value=torch.tensor(0.0), - ), - caplog.at_level(logging.WARNING), - ): - trainer.training_step(trainer.model, inputs) - - assert not any("trajectories" in r.message.lower() for r in caplog.records)