Bakery

bakery — where LLMs go to get baked. Distills arbitrary prefix contexts — system prompts, few-shot examples, conversation histories, accumulated memories — into model weights via KL-divergence training with LoRA. Zero inference-time prompt cost for the baked behavior.

Stack: Python (transformers + PEFT). Public repo. Generalization of prompt-baking 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 full prefix context, generates reference behavior.
• Student (adapters enabled) sees no prefix (or only the last N messages), trained to match the teacher's output distribution.

Objective minimizes per-token KL between teacher and student logits on whichever message tokens you mark as targets (default: all assistant turns).

Data sources

dataset accepts a local JSON file or a HuggingFace dataset ID (auto-detected). Format determines training mode:

Data format	Training mode
Prompts only (list of strings)	On-the-fly trajectory generation from teacher
Paired data (prompt+response, chat messages)	Train directly on precomputed pairs

Context baking

Beyond a single system prompt, bakery supports arbitrary prefix contexts via ContextConfig:

prefix_messages:
  - {role: system, content: "You answer concisely."}
  - {role: user, content: "Example Q"}
  - {role: assistant, content: "Example A"}
student_retained_turns: 0    # 0 = pure baking; N>0 = last N kept at inference
target_roles: [assistant]    # roles that contribute to KL loss
target_content_pattern: "^Answer:"   # optional regex over content

Or load from a file: prefix_messages_file: "./prefixes/persona_A.yaml".

Config is flat YAML parsed by HfArgumentParser, so any TrainingArguments field works.

Sources

• README

See also

• prompt-baking — the earlier, single-system-prompt version.