feat: generalize prompt baking to arbitrary prefix contexts

.github/workflows/prek.yml

@@ -1,9 +1,9 @@
 name: Prek
 on:
   push:
-    branches: main
+    branches: [main]
   pull_request:
-    branches: "*"
+  workflow_dispatch:
 
 jobs:
   prek:

.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"

.gitignore

@@ -173,3 +173,8 @@ cython_debug/
 # PyPI configuration file
 .pypirc
 .worktrees
+
+# Local-only / generated artifacts
+benchmarks/
+k8s/
+unsloth_compiled_cache/

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,33 +81,84 @@ 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,
 )
 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) |

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]"

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?"

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'?"

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"