feat: generalize prompt baking to arbitrary prefix contexts

[marksverdhei]

Summary

Pivots bakery from baking a single system prompt to baking arbitrary prefix contexts — conversation histories, accumulated memories, few-shot examples — so it becomes a tool for continual learning / memory distillation, not just prompt compression.

The teacher/student asymmetry is no longer hardcoded as [system, user] vs [user]; the teacher sees an arbitrary prefix_messages list, and the student sees an optionally-trimmed version (student_retained_turns: int).

What changed

• New ContextConfig dataclass — flat YAML fields: prefix_messages, prefix_messages_file, student_retained_turns, target_roles, target_content_pattern.
• New src/bakery/masking.py — per-token target mask builder using longest-common-token-id-prefix alignment (robust to BOS injection + template quirks), with a bounded cache.
• Renamed PromptBakingTrainer → ContextBakingTrainer. Old name kept as deprecated alias with a DeprecationWarning.
• Unified compute_loss and prediction_step via a shared _build_batch, eliminating ~90 lines of duplication.
• data.py — new create_conversational_dataset, collator handles both legacy and new batch shapes, HF messages loader preserves multi-turn history.
• CLI wires ContextConfig through HfArgumentParser; auto-desugars deprecated system_prompt → prefix_messages=[{role: system, content: ...}] with a warning.
• Per-row prefix_messages dataset column overrides the global prefix.
• Examples: multi_turn_prefix.yaml, continual_memory.yaml, per_row_prefix.yaml, pattern_targets.yaml.
• README — new "Context baking" section with migration note.

Backward compatibility

• system_prompt: "..." still works — desugars to prefix_messages=[{role: system, content: ...}] with a DeprecationWarning.
• PromptBakingTrainer import still works — deprecated alias of ContextBakingTrainer.
• Existing examples/basic.yaml and examples/sft_dataset.yaml unchanged.

Test plan

• pytest tests/test_context_baking.py tests/test_trainer.py — 21 passed (CPU, tiny GPT-2 + LoRA)
• Mask unit tests cover: single-assistant, multi-turn, regex filter, prefix exclusion via target_min_msg_idx, cache keying
• Trainer integration tests cover: system_prompt back-compat, global multi-turn prefix, per-row override, student_retained_turns, multi-turn conversational batch, target_roles restriction, pattern-filtered targets
• GPU smoke test on a real model (Qwen-0.6B) to confirm loss decreases on examples/multi_turn_prefix.yaml
• Backward-compat smoke test: examples/basic.yaml and examples/sft_dataset.yaml produce identical behavior to main

🤖 Generated with Claude Code