从“解题”到“证明”——OpenAI 通用大模型如何摘取 IMO 2025 金牌

pip install transformers==4.42 torch==2.4 einops

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch, random, time
tok = AutoTokenizer.from_pretrained("gpt2-medium")
model = AutoModelForCausalLM.from_pretrained("gpt2-medium")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def verdict(proof: str, question: str) -> float:
    """
    极简启发：检查是否出现关键词 QED / 矛盾 / 结论一致
    真实系统用 7B 验证模型 + 自洽采样
    """
    if "Q.E.D" in proof and proof.count("=>") >= 3:
        return 1.0
    if "contradiction" in proof:
        return 0.2            # 低分但不杀
    return 0.0

@torch.no_grad()
def rollout(prompt, max_new=400, temp=0.8):
    inputs = tok(prompt, return_tensors="pt").to(device)
    out = model.generate(**inputs, max_new_tokens=max_new, do_sample=True,
                         temperature=temp, pad_token_id=tok.eos_token_id)
    return tok.decode(out[0], skip_special_tokens=True)

def rstar(question, iterations=6, beam=3):
    pool = [question]                       # 初始 prompt
    for t in range(iterations):
        candidates = []
        for p in pool:
            for _ in range(beam):
                candidates.append(rollout(p, temp=0.5+random.random()*0.5))
        scored = [(c, verdict(c, question)) for c in candidates]
        scored.sort(key=lambda x: x[1], reverse=True)
        best = scored[:beam]
        pool = [b[0] for b in best]         # 下一轮递归输入
        print(f"step {t} best score = {best[0][1]:.2f}")
    return best[0][0]

if __name__ == "__main__":
    question = """
Let a, b, c > 0 such that a+b+c=1. Prove that
(a^2+b^2+c^2)^2 >= 3(a^3+b^3+c^3).
"""
    proof = rstar(question, iterations=5)
    print("\n=== Final Proof ===")
    print(proof)