推理模型从零构建 第5章：自我优化 (Self-Refinement)

import math

def heuristic_score(answer, prompt=None, brevity_bonus=500.0,
                     boxed_bonus=2.0, extract_bonus=1.0):
    score = 0.0

    # 有 \boxed{} 答案框的回答加分（说明格式规范）
    cand = extract_final_candidate(answer, fallback="none")
    if cand:
        score += boxed_bonus  # +2.0
    else:
        # 至少有一个数字也加分
        cand = extract_final_candidate(answer, fallback="number_only")
        if cand:
            score += extract_bonus  # +1.0

    # 简洁性奖励：回答越短得分越高（指数衰减）
    score += 1.5 * math.exp(-len(answer) / brevity_bonus)
    return score

# 两个同样正确的回答（x=83），但长度不同：
# Response 1 (1419字符): score = 2.088
# Response 2 (533字符):  score = 2.517 ← 更简洁，得分更高

# "Berlin" 的联合概率：5.9e-08（极小，难以比较）
# "Bridge" 的联合概率：1.0e-13

# 取 log 后：
# "Berlin" 的 joint log-prob: -16.6（可以直接比较大小）
# "Bridge" 的 joint log-prob: -29.9
# "Hamburg" 的 joint log-prob: -18.4
# → Berlin > Hamburg > Bridge，符合直觉

@torch.inference_mode()
def avg_logprob_answer(model, tokenizer, prompt, answer, device="cpu"):
    """计算模型对回答部分的平均 log-probability"""
    # 分别编码 prompt 和 answer
    prompt_ids = tokenizer.encode(prompt)
    answer_ids = tokenizer.encode(answer)
    full_ids = torch.tensor(prompt_ids + answer_ids, device=device)

    # 前向传播，得到所有位置的 log-probabilities
    logits = model(full_ids.unsqueeze(0)).squeeze(0)
    logprobs = torch.log_softmax(logits, dim=-1)

    # 只取回答部分的 token 得分
    start = len(prompt_ids) - 1
    end = full_ids.shape[0] - 1
    t_idx = torch.arange(start, end, device=device)
    next_tokens = full_ids[start + 1 : end + 1]
    next_token_logps = logprobs[t_idx, next_tokens]

    # 取平均（归一化掉长度影响）
    return torch.mean(next_token_logps)


# 测试：模型对正确答案的置信度更高
avg_logprob_answer(model, tokenizer,
    prompt="What is the capital of Germany?",
    answer=" The capital of Germany is Berlin.")  # → -0.204

avg_logprob_answer(model, tokenizer,
    prompt="What is the capital of Germany?",
    answer=" The capital of Germany is Bridge.")  # → -3.891

def make_critique_prompt(raw_prompt, draft):
    return (
        "You are a meticulous reviewer. Identify logical errors, "
        "missing steps, or arithmetic mistakes. If the answer seems "
        "correct, say so briefly. Then propose a concise plan to fix issues.\n\n"
        f"Question:\n{raw_prompt}\n\n"
        f"Draft answer:\n{draft}\n\n"
        "Write a short critique and bullet-point fix plan "
        "(under ~120 words).\nCritique:"
    )

def make_refine_prompt(raw_prompt, draft, critique):
    return (
        "Revise the answer using the critique. Keep it concise and "
        "end with a final boxed result: \\boxed{ANSWER}\n\n"
        f"Question:\n{raw_prompt}\n\n"
        f"Previous answer:\n{draft}\n\n"
        f"Critique:\n{critique}\n\n"
        "Revised answer:"
    )

def self_refinement_loop(model, tokenizer, raw_prompt, device,
                          iterations=2, score_fn=None, ...):
    # 生成初始回答
    current_full = generate(model, tokenizer, prompt, ...)
    current_score = score_fn(answer=current_full, prompt=prompt)

    for it in range(iterations):
        # 1. 模型自我批评
        critique = generate(model, tokenizer,
            make_critique_prompt(raw_prompt, current_full), ...)

        # 2. 根据批评修正
        revised = generate(model, tokenizer,
            make_refine_prompt(raw_prompt, current_full, critique), ...)
        revised_score = score_fn(answer=revised, prompt=prompt)

        # 3. 只在修正版更好时接受
        if revised_score >= current_score:
            current_full = revised
            current_score = revised_score

    return current_full


# 实际运行结果（seed=1, 2轮迭代）：
# 初始回答: \boxed{18}  (错误)   score = -0.855
# 第1轮修正: \boxed{83}  (正确!)  score = -0.226 → 接受
# 第2轮修正: \boxed{83}  (正确)   score = -1.320 → 拒绝（更差了）
# 最终答案: 83 ✓