"""5.1 图“同一步的概率分布在几个温度下”的数据：取一个提示第 1 步的整张词表 logits，按几个温度做 softmax。

用法：
    python temperature_step1.py                       # 默认提示“法国的首都是哪座城市？用一句话回答。”
    python temperature_step1.py --temps 0.5 1.0 1.5 --top 5

模型加载方式和 logprobs.py 相同（bf16，关掉思考模式）。输出存到 sample_output/temperature_step1_<设备>.json：
前几名候选的原始 logits，以及每个温度下它们的概率和“其余全部”的概率。
"""
import argparse
import json
from datetime import datetime
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from logprobs import DEFAULT_PROMPTS, pick_device, show


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='Qwen/Qwen3.5-2B')
    ap.add_argument('--prompt', default=DEFAULT_PROMPTS[0])
    ap.add_argument('--temps', type=float, nargs='+', default=[0.5, 1.0, 1.5])
    ap.add_argument('--top', type=int, default=5, help='单独列出前几名候选')
    ap.add_argument('--out', default='sample_output')
    args = ap.parse_args()

    device = pick_device()
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.float32 if device == 'cpu' else torch.bfloat16).to(device)
    model.eval()
    msgs = [{'role': 'user', 'content': args.prompt}]
    text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    inputs = tok(text, return_tensors='pt').to(model.device)
    with torch.no_grad():
        logits = model(**inputs).logits[0, -1].float().cpu()          # 第 1 步：整张词表的原始分数

    order = torch.argsort(logits, descending=True)[:args.top]
    rec = dict(model=args.model, device=device, date=f'{datetime.now():%Y-%m}', prompt=args.prompt, vocab=int(logits.shape[0]),
               top=[dict(token=show(tok, [int(i)]), id=int(i), logit=float(logits[i])) for i in order], temps={})
    for t in args.temps:
        p = (logits / t).softmax(-1)
        rec['temps'][str(t)] = dict(probs=[float(p[i]) for i in order], rest=float(1 - p[order].sum()))
    print(f'模型 {args.model}，设备 {device}，提示：{args.prompt}')
    print('候选 | logit | ' + ' | '.join(f'温度 {t}' for t in args.temps))
    for k, c in enumerate(rec['top']):
        print(f'{c["token"]} | {c["logit"]:.2f} | ' + ' | '.join(f'{rec["temps"][str(t)]["probs"][k]:.3f}' for t in args.temps))
    print('其余全部 | | ' + ' | '.join(f'{rec["temps"][str(t)]["rest"]:.3f}' for t in args.temps))
    Path(args.out).mkdir(exist_ok=True)
    path = Path(args.out) / f'temperature_step1_{device}.json'
    path.write_text(json.dumps(rec, ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'已存到 {path}')


if __name__ == '__main__':
    main()