"""层 1 关账判据 3：训练后 checkpoint 能被 from_pretrained 加载并生成通顺解答。

远程运行（CPU 即可，0.6B 生成 512 token 约 1-2 分钟）：
    python -u scripts/diag_generate.py
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from ars_opd.data import load_sft_dataset

MODEL_DIR = "/data/zym/outputs/sft_qwen3-0.6b_dapo1k"  # 正式 1 epoch 的产物

# 不挂 teacher 解答（teacher_completions_path 缺省）：只取题目做推理输入
ds = load_sft_dataset(
    "data/dapo-math-17k-unique.parquet", subset_size=1000, seed=42
)

tok = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForCausalLM.from_pretrained(MODEL_DIR, dtype=torch.float32)
model.eval()

# 取子集第 900+ 行附近的题（训练时见过，此处只验"会不会说话"不验泛化）
for i in (900, 950):
    prompt = tok.apply_chat_template(
        ds[i]["messages"],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,  # 必须与训练取值一致（docs/02 §2.3 边界契约）
    )
    inputs = tok(prompt, return_tensors="pt", add_special_tokens=False)
    with torch.no_grad():
        out = model.generate(
            **inputs, max_new_tokens=512, do_sample=False, temperature=None, top_p=None
        )
    completion = tok.decode(
        out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
    )
    print(f"===== 样本 {i} 题目 =====")
    print(ds[i]["messages"][-1]["content"][120:280], "…")
    print("----- 生成（前 600 字符）-----")
    print(completion[:600])
    print()

print(
    "判读：应为步骤化数学解答（markdown 风格、以 Answer: 行收尾的倾向）；"
    "乱码/复读/空输出 = 不通过。",
    flush=True,
)