Prompt Distillation

import asyncio
import mint
from mint import types

async def prompt_distillation():
    service_client = mint.ServiceClient()
    
    print("=== Prompt Distillation Pipeline ===")
    
    # Step 1: Use a larger teacher model to generate reasoning
    print("\n1. Sampling reasoning from teacher model...")
    
    teacher_sampler = service_client.create_sampling_client(
        base_model="Qwen/Qwen3-30B-A3B",  # Larger model
    ).result()
    tokenizer_teacher = teacher_sampler.get_tokenizer()
    renderer_teacher = mint.renderers.get_renderer("qwen3", tokenizer_teacher)
    
    # Sample from teacher on training tasks
    tasks = [
        "If x + 2 = 5, what is x?",
        "What is the capital of Japan?",
        "Explain why water boils at 100°C.",
    ]
    
    distillation_data = []
    
    for task in tasks:
        messages = [{"role": "user", "content": task}]
        prompt = renderer_teacher.build_generation_prompt(messages)
        
        sampling_params = types.SamplingParams(
            max_tokens=256,
            temperature=0.7,  # Some diversity
            stop=renderer_teacher.get_stop_sequences(),
        )
        
        output = teacher_sampler.sample(
            prompt, sampling_params=sampling_params, num_samples=1
        ).result()
        
        teacher_response = tokenizer_teacher.decode(output.sequences[0].tokens)
        
        distillation_data.append({
            "prompt": task,
            "reasoning": teacher_response,
        })
        print(f"  Task: {task[:30]}... -> {teacher_response[:50]}...")
    
    # Step 2: Fine-tune a smaller student model on teacher outputs
    print("\n2. Fine-tuning student model on teacher reasoning...")
    
    student_client = await service_client.create_lora_training_client_async(
        base_model="Qwen/Qwen3-0.6B",  # Smaller student
        rank=16,
    )
    tokenizer_student = student_client.get_tokenizer()
    adam_params = types.AdamParams(learning_rate=5e-5)
    
    for epoch in range(3):
        epoch_losses = []
        
        for example in distillation_data:
            prompt_text = example["prompt"]
            reasoning_text = example["reasoning"]
            
            # Tokenize prompt + reasoning
            full_text = f"{prompt_text} {reasoning_text}"
            tokens = tokenizer_student.encode(full_text)
            
            model_input = types.ModelInput.from_ints(tokens[:-1])
            target_tokens = tokens[1:]
            weights = [1.0] * len(target_tokens)
            
            datum = types.Datum(
                model_input=model_input,
                loss_fn_inputs={
                    "target_tokens": target_tokens,
                    "weights": weights,
                },
            )
            
            # Train student on teacher's output
            fb_future = student_client.forward_backward_async(
                [datum], loss_fn="cross_entropy"
            )
            result = await fb_future.result_async()
            epoch_losses.append(result.loss)
        
        optim_future = student_client.optim_step_async(adam_params)
        await optim_future.result_async()
        
        avg_loss = sum(epoch_losses) / len(epoch_losses)
        print(f"  Epoch {epoch}: loss={avg_loss:.4f}")
    
    # Step 3: Evaluate student on new tasks
    print("\n3. Evaluating distilled student model...")
    
    student_sampler = service_client.create_sampling_client_from_checkpoint(
        checkpoint_name=await student_client.save_weights_for_sampler_async(
            name="student-distilled-v1"
        ).result()
    ).result()
    
    new_task = "What is 3 * 4?"
    prompt_ids = tokenizer_student.encode(new_task)
    prompt_input = types.ModelInput.from_ints(prompt_ids)
    
    output = student_sampler.sample(
        prompt_input,
        sampling_params=types.SamplingParams(max_tokens=64, temperature=0.0),
    ).result()
    
    student_response = tokenizer_student.decode(output.sequences[0].tokens)
    print(f"  Student response: {student_response}")

asyncio.run(prompt_distillation())