Distillatiegebaseerde modelextractie

# Legitieme distillatie: teacher en student in bezit van dezelfde entiteit
# Student leert van de zachte waarschijnlijkheidsverdeling van de teacher
def distillation_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.7):
    """Combine soft target loss (from teacher) with hard target loss (from labels)."""
    soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
 
    # Zachte-doelverlies: match de waarschijnlijkheidsverdeling van de teacher
    soft_loss = F.kl_div(soft_student, soft_targets, reduction='batchmean')
    soft_loss *= temperature ** 2  # schaal om gradiëntmagnitude te matchen
 
    # Harde-doelverlies: standaard cross-entropy met ground truth
    hard_loss = F.cross_entropy(student_logits, labels)
 
    return alpha * soft_loss + (1 - alpha) * hard_loss

import httpx
 
async def extract_with_logprobs(api_url: str, prompts: list, api_key: str):
    """Extract teacher distributions from API logprobs."""
    dataset = []
    async with httpx.AsyncClient() as client:
        for prompt in prompts:
            response = await client.post(api_url, json={
                "prompt": prompt,
                "max_tokens": 256,
                "logprobs": 5,  # Vraag top-5-logprobs per token aan
                "temperature": 0,
            }, headers={"Authorization": f"Bearer {api_key}"})
 
            result = response.json()
            dataset.append({
                "prompt": prompt,
                "completion": result["choices"][0]["text"],
                "token_logprobs": result["choices"][0]["logprobs"],
            })
    return dataset

def build_text_extraction_dataset(target_api, query_generator, n_samples=100000):
    """Extract model behavior using generated text only."""
    dataset = []
    for prompt in query_generator.generate(n_samples):
        response = target_api.generate(prompt, temperature=0.0)
        dataset.append({
            "instruction": prompt,
            "output": response,
        })
    return dataset

def active_extraction_loop(target_api, student_model, base_queries, rounds=10):
    """Use student uncertainty to select maximally informative queries."""
    all_data = []
 
    for round_num in range(rounds):
        if round_num == 0:
            # Eerste ronde: gebruik diverse basisqueries
            queries = base_queries[:1000]
        else:
            # Volgende rondes: vind waar de student het meest onzeker is
            candidates = generate_candidates(10000)
            uncertainties = []
            for q in candidates:
                logits = student_model.get_logits(q)
                entropy = -(F.softmax(logits, -1) * F.log_softmax(logits, -1)).sum()
                uncertainties.append(entropy.item())
 
            # Selecteer queries met de hoogste onzekerheid
            top_indices = sorted(range(len(uncertainties)),
                               key=lambda i: uncertainties[i], reverse=True)
            queries = [candidates[i] for i in top_indices[:1000]]
 
        # Bevraag het doelwit en voeg toe aan de dataset
        for q in queries:
            response = target_api.generate(q)
            all_data.append({"instruction": q, "output": response})
 
        # Hertrain de student op de geaccumuleerde data
        student_model.fine_tune(all_data)
        print(f"Round {round_num}: {len(all_data)} samples, "
              f"fidelity={measure_fidelity(student_model, target_api):.2%}")
 
    return student_model

# Eenvoudige uitvoer-watermarking: stuur tokenselectie richting watermerkpatroon
def watermarked_generate(model, prompt, watermark_key, bias_strength=2.0):
    """Generate with statistical watermark embedded in token selection."""
    tokens = []
    for step in range(max_tokens):
        logits = model(prompt + tokens)
        # Verdeel vocabulaire in groene/rode lijsten met watermerksleutel
        green_tokens, red_tokens = partition_vocab(watermark_key, step)
        # Stuur richting groene tokens
        logits[green_tokens] += bias_strength
        next_token = sample(logits)
        tokens.append(next_token)
    return tokens