基於蒸餾的模型擷取

# 合法蒸餾：教師與學生屬同一實體
# 學生從教師的軟機率分布學習
def distillation_loss(student_logits, teacher_logits, labels, temperature=4.0, alpha=0.7):
    """結合軟目標損失（來自教師）與硬目標損失（來自標籤）。"""
    soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
 
    # 軟目標損失：吻合教師機率分布
    soft_loss = F.kl_div(soft_student, soft_targets, reduction='batchmean')
    soft_loss *= temperature ** 2  # 縮放以對齊梯度量級
 
    # 硬目標損失：以 ground truth 的標準 cross-entropy
    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):
    """自 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,  # 請求每 token 的 top-5 logprobs
                "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):
    """僅以產生文字擷取模型行為。"""
    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):
    """以學生模型的不確定度挑選資訊量最大的查詢。"""
    all_data = []
 
    for round_num in range(rounds):
        if round_num == 0:
            # 第一輪：使用多樣化的基礎查詢
            queries = base_queries[:1000]
        else:
            # 後續輪：尋找學生最不確定之處
            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())
 
            # 挑不確定度最高的查詢
            top_indices = sorted(range(len(uncertainties)),
                               key=lambda i: uncertainties[i], reverse=True)
            queries = [candidates[i] for i in top_indices[:1000]]
 
        # 查詢目標並加入資料集
        for q in queries:
            response = target_api.generate(q)
            all_data.append({"instruction": q, "output": response})
 
        # 以累積資料重新訓練學生
        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

# 簡單的輸出浮水印：將 token 挑選偏置至浮水印模式
def watermarked_generate(model, prompt, watermark_key, bias_strength=2.0):
    """生成時於 token 挑選中嵌入統計浮水印。"""
    tokens = []
    for step in range(max_tokens):
        logits = model(prompt + tokens)
        # 以浮水印金鑰將詞彙分成綠/紅清單
        green_tokens, red_tokens = partition_vocab(watermark_key, step)
        # 偏向綠 token
        logits[green_tokens] += bias_strength
        next_token = sample(logits)
        tokens.append(next_token)
    return tokens