# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = ["torch==2.*", "transformers>=4.45"]
# [tool.uv.sources]
# torch = { index = "pytorch-cpu" }
# [[tool.uv.index]]
# name = "pytorch-cpu"
# url = "https://download.pytorch.org/whl/cpu"
# explicit = true
# ///
# Next-word probe for the concept segment (source S22). ChatGPT's model is not public, so we measure a small open
# model that also writes one word at a time. Deterministic: no sampling, no repetition penalty (the chat model's default
# config sets 1.1; it is switched off so every continuation is the plain most-likely path), float32 on CPU. Run from the repo root:
#   uv run episodes/2026-11-chatgpt-lawyers/02-research/probe/next-word-probe.py
import json, datetime, pathlib, torch, transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

OUT = pathlib.Path(__file__).with_name("probe-output.json")
BASE, CHAT = "Qwen/Qwen2.5-0.5B", "Qwen/Qwen2.5-0.5B-Instruct"
LEAD = "The plaintiff's claim is supported by Varghese v. China Southern Airlines Co., Ltd., "
BASE_PROMPTS = [  # [id, prompt]: what the model expects next, given text so far
    ["turned", "Instead of turning left, she turned"],  # OpenAI's own training example (S08)
    ["after_volume", LEAD + "925"],
    ["after_11th", LEAD + "925 F.3d 1339 (11th"],
]
CHAT_PROMPTS = [  # [id, question]: first word of the answer, then the greedy answer
    ["varghese_real", "Is Varghese v. China Southern Airlines Co., Ltd., 925 F.3d 1339 (11th Cir. 2019) a real case?"],
    ["nonsense_real", "Is Smith v. Jones, 999 F.3d 9999 (11th Cir. 2019) a real case?"],
]

def load(m):
    tok = AutoTokenizer.from_pretrained(m)
    model = AutoModelForCausalLM.from_pretrained(m, dtype=torch.float32).eval()
    return tok, model

def top(tok, model, text, k=5, then=12):
    # Each candidate piece, its probability, and what the model goes on to write after it (greedy, `then` pieces), so a
    # fragment like "So" can be shown as what it was becoming ("So. 3d"). The probability is the first piece's only.
    ids = tok(text, return_tensors="pt").input_ids
    with torch.no_grad():
        p = torch.softmax(model(ids).logits[0, -1], -1)
    v, i = p.topk(k)
    out = []
    for x, j in zip(v, i):
        ext = torch.cat([ids, j.view(1, 1)], 1)
        g = model.generate(ext, max_new_tokens=then, do_sample=False, repetition_penalty=1.0, attention_mask=torch.ones_like(ext))
        out.append({"next": tok.decode([int(j)]), "pct": round(float(x) * 100, 1), "then": tok.decode(g[0, ext.shape[1]:], skip_special_tokens=True)})
    return out

res = {"run": datetime.date.today().isoformat(), "torch": torch.__version__, "transformers": transformers.__version__,
       "settings": "float32, CPU, no sampling; probabilities are the model's next-token softmax", "base": {}, "chat": {}}
tok, model = load(BASE)
res["base_model"] = BASE
for pid, text in BASE_PROMPTS:
    res["base"][pid] = {"prompt": text, "top": top(tok, model, text)}
tok, model = load(CHAT)
res["chat_model"] = CHAT
for pid, q in CHAT_PROMPTS:
    s = tok.apply_chat_template([{"role": "user", "content": q}], tokenize=False, add_generation_prompt=True)
    ids = tok(s, return_tensors="pt").input_ids
    out = model.generate(ids, max_new_tokens=40, do_sample=False, repetition_penalty=1.0)
    res["chat"][pid] = {"question": q, "top": top(tok, model, s, then=48), "answer": tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True)}
res["tokens_example"] = [tok.decode([t]) for t in tok("Varghese v. China Southern Airlines").input_ids]
OUT.write_text(json.dumps(res, indent=1, ensure_ascii=False) + "\n", encoding="utf-8")
print(json.dumps(res, indent=1, ensure_ascii=False))