Distillation: Compressing Stable Behaviors

# optimization/collect_teacher_data.py
"""Collect teacher model outputs for distillation training data.

Runs the teacher model over a set of examples and filters for
correct outputs to use as training data.
"""

import json
from pathlib import Path


def load_examples(path: str) -> list[dict]:
    """Load examples from a JSONL file.

    Each line should have at minimum: question, and optionally
    conversation_context and expected_label for filtering.
    """
    examples = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                examples.append(json.loads(line))
    return examples


def format_routing_prompt(question: str, context: str = "") -> str:
    """Format a retrieval-routing classification prompt.

    This is the same prompt your pipeline uses for routing.
    """
    prompt = (
        "Classify this question into one of four retrieval modes.\n\n"
        "Modes:\n"
        "- code: question is about specific code, files, or symbols\n"
        "- docs: question is about documentation, README, or guides\n"
        "- hybrid: question spans both code and documentation\n"
        "- skip: question is general knowledge, no retrieval needed\n\n"
    )
    if context:
        prompt += f"Conversation context:\n{context}\n\n"
    prompt += f"Question: {question}\n\nMode:"
    return prompt


def collect_teacher_outputs(
    examples: list[dict],
    generate_fn,
) -> list[dict]:
    """Run the teacher model and collect outputs.

    Args:
        examples: List of dicts with 'question' and optional
            'conversation_context' and 'expected_label'.
        generate_fn: Function that takes a prompt string and
            returns the model's text output.

    Returns:
        List of training examples with teacher outputs.
    """
    results = []

    for ex in examples:
        prompt = format_routing_prompt(
            ex["question"],
            ex.get("conversation_context", ""),
        )
        teacher_output = generate_fn(prompt).strip().lower()

        result = {
            "prompt": prompt,
            "completion": teacher_output,
            "question": ex["question"],
        }

        # If we have expected labels, mark whether teacher was correct
        if "expected_label" in ex:
            result["expected"] = ex["expected_label"]
            result["teacher_correct"] = (
                teacher_output == ex["expected_label"]
            )

        results.append(result)

    return results


def filter_correct(results: list[dict]) -> list[dict]:
    """Keep only examples where the teacher produced the correct output."""
    return [r for r in results if r.get("teacher_correct", True)]


def save_training_data(results: list[dict], output_path: str) -> None:
    """Save filtered results as training data in JSONL format."""
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)

    with open(path, "w") as f:
        for r in results:
            training_example = {
                "prompt": r["prompt"],
                "completion": r["completion"],
            }
            f.write(json.dumps(training_example) + "\n")

    print(f"Saved {len(results)} training examples to {output_path}")

pip install peft trl unsloth transformers datasets

# optimization/train_student.py
"""Train a student model via QLoRA distillation.

Uses Unsloth + TRL for efficient parameter-efficient fine-tuning
on teacher-generated training data.
"""

from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig
from datasets import load_dataset


def format_for_training(example: dict) -> dict:
    """Format a prompt-completion pair for SFT training."""
    text = (
        f"### Instruction:\n{example['prompt']}\n\n"
        f"### Response:\n{example['completion']}"
    )
    return {"text": text}


def train(
    model_name: str = "unsloth/Qwen2.5-3B-bnb-4bit",
    training_data_path: str = "optimization/training_data/routing.jsonl",
    output_dir: str = "optimization/models/routing-student",
    max_steps: int = 200,
    learning_rate: float = 2e-4,
    lora_rank: int = 16,
):
    """Run QLoRA training on teacher-generated data.

    Args:
        model_name: Hugging Face model ID. Unsloth provides
            pre-quantized 4-bit versions for efficient training.
        training_data_path: Path to JSONL training data.
        output_dir: Where to save the trained LoRA adapter.
        max_steps: Training steps. Start small and increase if
            validation loss is still decreasing.
        learning_rate: Learning rate. 2e-4 is a good default for
            QLoRA; lower if training is unstable.
        lora_rank: LoRA rank. Higher rank = more capacity but more
            memory. 16 is a good starting point.
    """
    # Load the base model with 4-bit quantization
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=model_name,
        max_seq_length=2048,
        load_in_4bit=True,
    )

    # Add LoRA adapters to the model
    model = FastLanguageModel.get_peft_model(
        model,
        r=lora_rank,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        lora_alpha=lora_rank,  # Common default: alpha == rank
        lora_dropout=0,
        use_gradient_checkpointing="unsloth",
    )

    # Load and format training data
    dataset = load_dataset("json", data_files=training_data_path, split="train")
    dataset = dataset.map(format_for_training)

    # Configure training
    training_config = SFTConfig(
        output_dir=output_dir,
        max_steps=max_steps,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        learning_rate=learning_rate,
        logging_steps=10,
        save_steps=50,
        warmup_steps=20,
        fp16=True,
        dataset_text_field="text",
        max_seq_length=2048,
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        args=training_config,
    )

    # Train
    print(f"Training on {len(dataset)} examples for {max_steps} steps...")
    trainer.train()

    # Save the LoRA adapter (not the full model)
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    print(f"Saved LoRA adapter to {output_dir}")


if __name__ == "__main__":
    train()

python optimization/train_student.py

# optimization/eval_student.py
"""Evaluate student model against teacher on the benchmark.

Compares student accuracy, latency, and cost against the teacher
to determine whether distillation was worthwhile.
"""

import json
import time
from pathlib import Path

from unsloth import FastLanguageModel


def load_student(
    model_name: str = "unsloth/Qwen2.5-3B-bnb-4bit",
    adapter_path: str = "optimization/models/routing-student",
):
    """Load the base model with the trained LoRA adapter."""
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=model_name,
        max_seq_length=2048,
        load_in_4bit=True,
    )
    # Load LoRA weights on top of base model
    model.load_adapter(adapter_path)
    FastLanguageModel.for_inference(model)
    return model, tokenizer


def evaluate(
    model,
    tokenizer,
    eval_path: str = "benchmark/routing_examples.jsonl",
) -> dict:
    """Run student model on eval set and compute metrics."""
    examples = []
    with open(eval_path) as f:
        for line in f:
            line = line.strip()
            if line:
                examples.append(json.loads(line))

    correct = 0
    total = 0
    latencies = []
    errors = []

    for ex in examples:
        if "expected_label" not in ex:
            continue

        prompt = (
            f"### Instruction:\n"
            f"{format_routing_prompt(ex['question'], ex.get('conversation_context', ''))}\n\n"
            f"### Response:\n"
        )

        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        start = time.perf_counter()
        outputs = model.generate(
            **inputs,
            max_new_tokens=10,
            temperature=0,
        )
        elapsed = time.perf_counter() - start

        response = tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True,
        ).strip().lower()

        total += 1
        latencies.append(elapsed)

        if response == ex["expected_label"]:
            correct += 1
        else:
            errors.append({
                "question": ex["question"],
                "expected": ex["expected_label"],
                "got": response,
            })

    return {
        "accuracy": correct / total if total > 0 else 0,
        "total": total,
        "correct": correct,
        "avg_latency_ms": sum(latencies) / len(latencies) * 1000 if latencies else 0,
        "errors": errors[:10],  # First 10 errors for inspection
    }


def format_routing_prompt(question: str, context: str = "") -> str:
    """Same formatting as teacher data collection."""
    prompt = (
        "Classify this question into one of four retrieval modes.\n\n"
        "Modes:\n"
        "- code: question is about specific code, files, or symbols\n"
        "- docs: question is about documentation, README, or guides\n"
        "- hybrid: question spans both code and documentation\n"
        "- skip: question is general knowledge, no retrieval needed\n\n"
    )
    if context:
        prompt += f"Conversation context:\n{context}\n\n"
    prompt += f"Question: {question}\n\nMode:"
    return prompt


if __name__ == "__main__":
    print("Loading student model...")
    model, tokenizer = load_student()

    print("Evaluating...")
    results = evaluate(model, tokenizer)

    print(f"\nStudent Evaluation Results:")
    print(f"  Accuracy: {results['accuracy']:.1%}")
    print(f"  Correct: {results['correct']}/{results['total']}")
    print(f"  Avg latency: {results['avg_latency_ms']:.0f}ms")

    if results["errors"]:
        print(f"\nSample errors:")
        for err in results["errors"][:5]:
            print(f"  Q: {err['question'][:60]}...")
            print(f"    Expected: {err['expected']}, Got: {err['got']}")

python optimization/eval_student.py

Loading student model...
Evaluating...

Student Evaluation Results:
  Accuracy: 91.3%
  Correct: 42/46
  Avg latency: 45ms

Sample errors:
  Q: What does the authentication middleware do and where is i...
    Expected: hybrid, Got: code
  Q: How do I set up the project for the first time?...
    Expected: docs, Got: hybrid

# optimization/compare_models.py
"""Compare student and teacher performance side by side."""


def compare(teacher_results: dict, student_results: dict) -> None:
    """Print a comparison table."""
    print(f"\n{'Metric':<25} {'Teacher':>12} {'Student':>12} {'Delta':>12}")
    print("-" * 65)

    t_acc = teacher_results["accuracy"]
    s_acc = student_results["accuracy"]
    print(f"{'Accuracy':<25} {t_acc:>11.1%} {s_acc:>11.1%} {s_acc - t_acc:>+11.1%}")

    t_lat = teacher_results["avg_latency_ms"]
    s_lat = student_results["avg_latency_ms"]
    print(f"{'Avg latency (ms)':<25} {t_lat:>11.0f} {s_lat:>11.0f} {s_lat - t_lat:>+11.0f}")

    # Cost comparison depends on your setup
    # For hosted teacher vs local student, the student cost approaches zero
    print(f"\n{'Verdict':}")
    print("-" * 40)

    quality_drop = t_acc - s_acc
    speed_gain = t_lat / s_lat if s_lat > 0 else float("inf")

    if quality_drop <= 0.05 and speed_gain >= 2:
        print("  Student is viable: <5% quality drop with significant speed gain.")
    elif quality_drop <= 0.10:
        print("  Student is marginal: 5-10% quality drop. Consider more training data.")
    else:
        print("  Student is not ready: >10% quality drop. Needs more data or larger student.")