"""5.1 一个 token 是怎么来的：看模型每一步的候选 token 和概率，再在几个温度下各采样几次。

用法：
    python logprobs.py                                   # 两个默认提示，默认模型 Qwen/Qwen3.5-2B
    python logprobs.py --prompt "你自己的问题"
    python logprobs.py --model Qwen/Qwen3.5-0.8B         # 内存小的机器换小模型
    python logprobs.py --temps 0.3 1.0 1.5 --samples 5

在 NVIDIA 显卡上自动用 CUDA，在 Apple Silicon Mac 上用 MPS，都没有就用 CPU（慢，但能跑）。
"""
import argparse
import json
from datetime import datetime
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

DEFAULT_PROMPTS = [
    '法国的首都是哪座城市？用一句话回答。',
    '推荐一家上海的独立咖啡馆，给出店名和具体地址，一句话回答。',
]


def pick_device():
    if torch.cuda.is_available():
        return 'cuda'
    if torch.backends.mps.is_available():
        return 'mps'
    return 'cpu'


def show(tok, ids):
    """把一个 token 显示成可读的样子：空格显示成 ·，换行显示成 ⏎。"""
    return tok.decode(ids).replace(' ', '·').replace('\n', '⏎') or '∅'


def generate(model, tok, prompt, temperature, max_new_tokens, n=1, top_p=1.0):
    """生成 n 条回复。返回每条回复的 token id，以及每一步的原始 logits（未经温度调整的模型分布）。"""
    msgs = [{'role': 'user', 'content': 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)
    kw = dict(do_sample=False) if temperature == 0 else dict(do_sample=True, temperature=temperature, top_p=top_p, top_k=0)
    out = model.generate(**inputs, max_new_tokens=max_new_tokens, num_return_sequences=n, output_logits=True,
                         return_dict_in_generate=True, pad_token_id=tok.eos_token_id, **kw)
    seqs = out.sequences[:, inputs['input_ids'].shape[1]:]
    probs = torch.stack(out.logits, dim=1).float().softmax(-1)          # [n, 步数, 词表]
    return seqs, probs


def step_table(tok, seq, probs, low, top_k=5, top_p=0.9):
    """贪心解码的逐步表：选中的 token、它的概率、其他候选，以及累计概率到 top_p 需要几个候选。"""
    rows = []
    for i, tid in enumerate(seq.tolist()):
        if tid == tok.eos_token_id:
            break
        p = probs[i]
        top = torch.topk(p, top_k)
        nucleus = int((torch.cumsum(torch.sort(p, descending=True).values, 0) < top_p).sum()) + 1
        rows.append(dict(step=i + 1, token=show(tok, [tid]), prob=float(p[tid]), nucleus=nucleus,
                         alternatives=[(show(tok, [int(j)]), float(q)) for q, j in zip(top.values, top.indices) if int(j) != tid][:top_k - 1]))
    print('| 步 | 选中的 token | 概率 | 其他候选（概率） | 凑满 90% 概率要几个候选 |')
    print('|---|---|---|---|---|')
    for r in rows:
        flag = ' ⚠' if r['prob'] < low else ''
        alts = '，'.join(f'{t} {q:.2f}' for t, q in r['alternatives'])
        print(f'| {r["step"]} | {r["token"]} | {r["prob"]:.2f}{flag} | {alts} | {r["nucleus"]} |')
    return rows


def marked(tok, seq, probs, low):
    """把回复里概率低于 low 的 token 用【】标出来，后面跟上它的概率。"""
    parts = []
    for i, tid in enumerate(seq.tolist()):
        if tid == tok.eos_token_id:
            break
        s = tok.decode([tid])
        p = float(probs[i][tid])
        parts.append(f'【{s}|{p:.2f}】' if p < low else s)
    return ''.join(parts).replace('\n', ' ')


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='Qwen/Qwen3.5-2B')
    ap.add_argument('--prompt', action='append', help='可以给多次；不给就用两个默认提示')
    ap.add_argument('--temps', type=float, nargs='+', default=[0.3, 0.7, 1.0, 1.5])
    ap.add_argument('--samples', type=int, default=5, help='每个温度采样几次')
    ap.add_argument('--top-p', type=float, default=0.9, help='最后在温度 --top-p-temp 下加上 top-p 再采样一轮，对比效果')
    ap.add_argument('--top-p-temp', type=float, default=1.0)
    ap.add_argument('--max-new-tokens', type=int, default=48)
    ap.add_argument('--low', type=float, default=0.3, help='概率低于它的 token 会被标出来')
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--out', default='results')
    args = ap.parse_args()

    device = pick_device()
    torch.manual_seed(args.seed)
    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()
    print(f'模型 {args.model}，设备 {device}，{datetime.now():%Y-%m-%d %H:%M}。')
    print(f'概率一律指模型原始分布（温度 1）下的概率；低于 {args.low} 的 token 标 ⚠，或在回复里写成【token|概率】。')
    print('采样时不做 top-k 截断；除最后一组外 top-p 都是 1，即只有温度在起作用。')

    record = dict(model=args.model, device=device, prompts=[])
    for prompt in args.prompt or DEFAULT_PROMPTS:
        print(f'\n# 提示：{prompt}\n')
        print('## 温度 0（每步都选概率最高的那个）\n')
        seqs, probs = generate(model, tok, prompt, 0, args.max_new_tokens)
        rows = step_table(tok, seqs[0], probs[0], args.low)
        print(f'\n回复：{marked(tok, seqs[0], probs[0], args.low)}')
        rec = dict(prompt=prompt, greedy=rows, samples={})
        settings = [(t, 1.0) for t in args.temps] + ([(args.top_p_temp, args.top_p)] if args.top_p < 1 else [])
        for t, top_p in settings:
            label = f'温度 {t}' + (f' + top-p {top_p}' if top_p < 1 else '')
            print(f'\n## {label}，采样 {args.samples} 次\n')
            seqs, probs = generate(model, tok, prompt, t, args.max_new_tokens, n=args.samples, top_p=top_p)
            texts = []
            for k in range(args.samples):
                m = marked(tok, seqs[k], probs[k], args.low)
                n_low = m.count('【')
                print(f'{k + 1}. {m}（低概率 token {n_low} 个）')
                texts.append(dict(text=tok.decode(seqs[k], skip_special_tokens=True), marked=m, n_low=n_low))
            print(f'\n{args.samples} 次里有 {len({x["text"] for x in texts})} 种不同的回复。')
            rec['samples'][label] = texts
        record['prompts'].append(rec)

    Path(args.out).mkdir(exist_ok=True)
    path = Path(args.out) / f'logprobs_{args.model.split("/")[-1]}_{device}_{datetime.now():%Y%m%d_%H%M}.json'
    path.write_text(json.dumps(record, ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'\n完整结果已存到 {path}')


if __name__ == '__main__':
    main()