#!/usr/bin/env python3
"""T1.4 FP8 acceptance check on the exact generation config (DECISIONS.md S-20).

Spec §6 criteria: teacher-forced KL(bf16 || FP8) <= 0.02 nats/token, and GSM8K-250 accuracy within
1.5 points of bf16. "FP8" is the whole generation config in configs/gen/teacher_fp8.yaml — FP8 block
weights AND FP8 KV cache, router bf16, Triton MoE and attention — run on gen's own kernels. The bf16
reference is configs/gen/teacher_bf16_ref.yaml (same kernels, bf16 weights and KV, part of the
weights in pinned CPU memory).

KL is exact over the full 151,936-token vocabulary: the capture step records the final hidden state
at every position (the input of vLLM's compute_logits), and the kl step rebuilds both models' logits
through the shared, unquantised bf16 LM head in strict fp32. Triton attention reads K/V from the
(FP8) cache in prefill as in decode, so teacher-forced positions carry the same KV quantisation as
generation.

Subcommands: prep-gsm8k, prep-math, prep-tf, generate, seqs-from, capture, kl, score-gsm8k,
analyze-trunc. Every vLLM step runs as its own process through the gen queue.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import urllib.request
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np
import yaml

from lowbit.eval.gsm8k import MATH_SUFFIX, extract_answer, same_number  # noqa: F401 (re-exported)
from lowbit.gen.loops import loop_score  # noqa: E402,F401 (re-exported for router_sampling_sweep.py)

REPO = Path(__file__).resolve().parent.parent
DATA = Path(os.environ.get("MT_DATA", "/path/to/data"))
EOS_DOC = 151643  # <|endoftext|>, the document separator in packed text
PROD_CFG = REPO / "configs" / "gen" / "teacher_fp8.yaml"  # spec §6 sampling lives here
PARQUET = "https://huggingface.co/datasets/{ds}/resolve/refs%2Fconvert%2Fparquet/{cfg}/{split}/0000.parquet"
MARCO_AYA = {"en": "English", "zh": "Simplified Chinese", "ar": "Standard Arabic", "de": "German", "es": "Spanish",
             "fr": "French", "ko": "Korean", "ja": "Japanese", "pt": "Portuguese", "tr": "Turkish",
             "id": "Indonesian", "it": "Italian", "nl": "Dutch", "pl": "Polish", "ru": "Russian",
             "vi": "Vietnamese", "th": "Thai", "uk": "Ukrainian", "ms": "Standard Malay", "bn": "Bengali",
             "ur": "Urdu", "el": "Greek", "hu": "Hungarian", "ne": "Nepali"}  # Aya lacks he, cs, kk, ro, az


def fetch(url: str, dest: Path) -> str:
    dest.parent.mkdir(parents=True, exist_ok=True)
    if not dest.exists():
        urllib.request.urlretrieve(url, dest)
    return hashlib.sha256(dest.read_bytes()).hexdigest()


def write_jsonl(path: Path, rows) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def read_jsonl(path) -> list[dict]:
    return [json.loads(l) for l in open(path) if l.strip()]


def load_tokenizer(model: str):
    from transformers import AutoTokenizer
    return AutoTokenizer.from_pretrained(model)


# ---------------------------------------------------------------- data prep (CPU)


def prep_gsm8k(a) -> None:
    import pyarrow.parquet as pq
    dest = DATA / "evalsets/gsm8k/main_test.parquet"
    sha = fetch(PARQUET.format(ds="openai/gsm8k", cfg="main", split="test"), dest)
    rows = pq.read_table(dest).to_pylist()
    ids = sorted(random.Random(a.seed).sample(range(len(rows)), 250))
    subset = {"dataset": "openai/gsm8k", "config": "main", "split": "test", "dataset_sha": a.dataset_sha,
              "parquet_sha256": sha, "seed": a.seed, "n": 250, "ids": ids}
    out = REPO / "configs/eval/subsets/gsm8k_250.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(subset, indent=1) + "\n")
    write_jsonl(DATA / "fp8check/gsm8k_250.jsonl",
                [{"id": f"gsm8k-{i}", "text": rows[i]["question"] + MATH_SUFFIX,
                  "gold": rows[i]["answer"].split("####")[-1].strip().replace(",", ""), "kind": "gsm8k"}
                 for i in ids])
    print(f"gsm8k: {len(rows)} test items, 250 sampled (seed {a.seed}) -> {out}")


def prep_math(a) -> None:
    import pyarrow.parquet as pq
    subjects = ["algebra", "counting_and_probability", "geometry", "intermediate_algebra", "number_theory",
                "prealgebra", "precalculus"]
    pool = []
    for s in subjects:
        dest = DATA / f"evalsets/hendrycks_math/{s}_train.parquet"
        fetch(PARQUET.format(ds="EleutherAI/hendrycks_math", cfg=s, split="train"), dest)
        pool += [dict(r, subject=s) for r in pq.read_table(dest).to_pylist() if r["level"] in ("Level 4", "Level 5")]
    pick = random.Random(a.seed).sample(pool, a.n)
    write_jsonl(DATA / "fp8check/math_long.jsonl",
                [{"id": f"math-{i}", "text": r["problem"] + MATH_SUFFIX, "kind": "math", "level": r["level"],
                  "subject": r["subject"]} for i, r in enumerate(pick)])
    print(f"math: {len(pool)} level-4/5 train problems, {a.n} sampled")


def prep_tf(a) -> None:
    """Aya human-written prompt+completion text in Marco's languages, packed into fixed windows."""
    import pyarrow.parquet as pq
    tok = load_tokenizer(a.model)
    t = pq.read_table(DATA / "bench/aya_train.parquet", columns=["inputs", "targets", "language"]).to_pylist()
    rng = random.Random(a.seed)
    by_lang = defaultdict(list)
    names = {v: k for k, v in MARCO_AYA.items()}
    for r in t:
        if r["language"] in names:
            by_lang[names[r["language"]]].append(r["inputs"].strip() + "\n\n" + r["targets"].strip())
    per_lang = a.tokens // len(by_lang)
    docs = []
    for lang, texts in sorted(by_lang.items()):
        rng.shuffle(texts)
        n = 0
        for txt in texts:
            ids = tok(txt, add_special_tokens=False)["input_ids"]
            docs.append((lang, ids))
            n += len(ids) + 1
            if n >= per_lang:
                break
    rng.shuffle(docs)
    windows, cur, segs = [], [], []
    for lang, ids in docs:
        ids = ids + [EOS_DOC]
        while ids:
            room = a.window - len(cur)
            segs.append([len(cur), lang])
            cur += ids[:room]
            ids = ids[room:]
            if len(cur) == a.window:
                windows.append({"id": f"aya-{len(windows)}", "kind": "aya", "token_ids": cur, "segments": segs})
                cur, segs = [], []
    write_jsonl(DATA / "fp8check/tf_aya.jsonl", windows)
    langs = Counter()
    for w in windows:
        bounds = [s[0] for s in w["segments"]] + [a.window]
        for (start, lang), end in zip(w["segments"], bounds[1:]):
            langs[lang] += end - start
    print(f"prep-tf: {len(windows)} windows x {a.window} = {len(windows) * a.window:,} tokens; per language: {dict(langs)}")


# ---------------------------------------------------------------- vLLM steps (GPU, via the gen queue)


def build_llm(cfg_path: str, overrides: list[str], in_process: bool):
    cfg = yaml.safe_load(open(cfg_path))
    os.environ.update({k: str(v) for k, v in cfg.get("env", {}).items()})
    if in_process:  # the capture hook must live in the process that runs the model
        os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
    from vllm import LLM
    eng = dict(cfg["engine"])
    for kv in overrides:
        k, v = kv.split("=", 1)
        eng[k] = yaml.safe_load(v)
    return LLM(model=cfg["model"], seed=0, **eng), cfg


def spec_sampling(path: str | Path = PROD_CFG) -> dict:
    """Spec §6 sampling, always from the production generation config: an FP8-vs-bf16 comparison
    must sample both engines identically, and the bf16 reference config carries no sampling block."""
    return yaml.safe_load(open(path))["sampling"]


def generate(a) -> None:
    llm, _ = build_llm(a.engine, a.override, in_process=False)
    from vllm import SamplingParams
    from vllm.inputs import TokensPrompt
    tok = llm.get_tokenizer()
    rows = read_jsonl(a.prompts)[: a.limit or None]
    prompts = [TokensPrompt(prompt_token_ids=tok(tok.apply_chat_template([{"role": "user", "content": r["text"]}],
                                                                       add_generation_prompt=True, tokenize=False),
                                                 add_special_tokens=False)["input_ids"]) for r in rows]
    if a.sampling == "greedy":
        sp = SamplingParams(temperature=0.0, max_tokens=a.max_tokens)
    else:
        s = spec_sampling(a.sampling_config)
        sp = SamplingParams(temperature=s["temperature"], top_p=s["top_p"], max_tokens=a.max_tokens, seed=0)
    outs = llm.generate(prompts, sp)
    write_jsonl(Path(a.out), [{"id": r.get("id", f"p{i}"), **{k: v for k, v in r.items() if k != "text"},
                               "prompt_token_ids": o.prompt_token_ids,
                               "output_token_ids": list(o.outputs[0].token_ids), "text": o.outputs[0].text,
                               "finish_reason": o.outputs[0].finish_reason} for i, (r, o) in enumerate(zip(rows, outs))])
    n_tok = sum(len(o.outputs[0].token_ids) for o in outs)
    print(json.dumps({"n": len(outs), "out_tokens": n_tok,
                      "truncated": sum(o.outputs[0].finish_reason == "length" for o in outs)}))


def seqs_from(a) -> None:
    """Teacher-forcing sequences from generate outputs: chat prompt + response, response start marked."""
    rows = read_jsonl(a.generated)
    write_jsonl(Path(a.out), [{"id": r["id"], "kind": r.get("kind", "gen"),
                               "token_ids": r["prompt_token_ids"] + r["output_token_ids"],
                               "response_start": len(r["prompt_token_ids"]), "finish_reason": r["finish_reason"]}
                              for r in rows if r["output_token_ids"]])


def install_capture(store: dict):
    """Wrap the V2 runner's prompt-logprob step to keep each prefill row's final hidden state."""
    from vllm.v1.worker.gpu.sample import prompt_logprob as pl
    orig = pl.PromptLogprobsWorker.compute_prompt_logprobs

    def wrapped(self, logits_fn, hidden_states, input_batch, all_token_ids, num_computed_tokens, prompt_lens):
        plens = prompt_lens[input_batch.idx_mapping_np]
        computed = input_batch.num_computed_prefill_tokens_np
        qsl = input_batch.query_start_loc_np
        for i, rid in enumerate(input_batch.req_ids):
            if computed[i] >= plens[i]:
                continue  # decode row
            s, e = int(qsl[i]), int(qsl[i + 1])
            store[rid].append((int(computed[i]), hidden_states[s:e].detach().to("cpu", copy=True)))
        return orig(self, logits_fn, hidden_states, input_batch, all_token_ids, num_computed_tokens, prompt_lens)

    pl.PromptLogprobsWorker.compute_prompt_logprobs = wrapped


def capture(a) -> None:
    store: dict[str, list] = defaultdict(list)
    llm, _ = build_llm(a.engine, a.override, in_process=True)
    install_capture(store)  # after construction: the warm-up runs carry no prompt logprobs
    import torch
    from vllm import SamplingParams
    from vllm.inputs import TokensPrompt
    seqs = []
    for f in a.sequences:
        seqs += read_jsonl(f)
    seqs = seqs[: a.limit or None]
    sp = SamplingParams(temperature=0.0, max_tokens=1, prompt_logprobs=0)
    outs = llm.generate([TokensPrompt(prompt_token_ids=s["token_ids"]) for s in seqs], sp)
    out_dir = Path(a.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    total = sum(len(s["token_ids"]) for s in seqs)
    hid = np.lib.format.open_memmap(out_dir / "hidden_bf16.npy", mode="w+", dtype=np.uint16, shape=(total, 1024))
    vllm_lp = np.full(total, np.nan, dtype=np.float32)  # vLLM's own logprob of token p given < p
    index, off = [], 0
    for s, o in zip(seqs, outs):
        rid = o.request_id
        chunks = store.get(rid) or next((v for k, v in store.items() if k.startswith(rid + "-") or rid.startswith(k)), None)
        if chunks is None:
            raise RuntimeError(f"no captured hidden states for request {rid} ({s['id']}); ids seen: {list(store)[:5]}")
        chunks.sort(key=lambda c: c[0])
        h = torch.cat([c[1] for c in chunks])
        n = len(s["token_ids"])
        starts = [c[0] for c in chunks]
        if h.shape[0] != n or starts[0] != 0:
            raise RuntimeError(f"{s['id']}: captured {h.shape[0]} rows from {starts}, expected {n}")
        hid[off:off + n] = h.to(torch.bfloat16).view(torch.int16).numpy().view(np.uint16)
        for p, d in enumerate(o.prompt_logprobs or []):
            if d is not None:
                vllm_lp[off + p] = d[s["token_ids"][p]].logprob
        index.append({k: v for k, v in s.items() if k != "token_ids"} | {"offset": off, "length": n})
        off += n
    hid.flush()
    np.save(out_dir / "vllm_token_logprob.npy", vllm_lp)
    write_jsonl(out_dir / "index.jsonl", index)
    (out_dir / "meta.json").write_text(json.dumps({"engine": a.engine, "override": a.override, "positions": total,
                                                   "sequences": len(seqs)}, indent=1))
    print(json.dumps({"captured_positions": total, "sequences": len(seqs)}))


def kl(a) -> None:
    import torch
    from safetensors import safe_open
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.set_float32_matmul_precision("highest")
    emb = None
    for f in sorted(Path(a.model).glob("*.safetensors")):
        with safe_open(f, "pt") as st:
            if "model.embed_tokens.weight" in st.keys():
                emb = st.get_tensor("model.embed_tokens.weight")
    dev = a.device
    E = emb.to(dev, torch.float32)  # tied LM head, bf16 -> fp32 exactly
    ref_dir, test_dir = Path(a.ref), Path(a.test)
    ref_idx = {r["id"]: r for r in read_jsonl(ref_dir / "index.jsonl")}
    test_idx = {r["id"]: r for r in read_jsonl(test_dir / "index.jsonl")}
    seqs = {}
    for f in a.sequences:
        seqs.update({s["id"]: s for s in read_jsonl(f)})
    H_ref = np.load(ref_dir / "hidden_bf16.npy", mmap_mode="r")
    H_test = np.load(test_dir / "hidden_bf16.npy", mmap_mode="r")
    lp_ref_vllm = np.load(ref_dir / "vllm_token_logprob.npy")
    cols = defaultdict(list)
    for sid, r in ref_idx.items():
        if sid not in test_idx:  # compare the positions both captures have (subset diagnostics)
            continue
        t, s = test_idx[sid], seqs[sid]
        n = r["length"]
        assert t["length"] == n and len(s["token_ids"]) == n, sid
        hb = torch.from_numpy(np.array(H_ref[r["offset"]:r["offset"] + n], copy=True).view(np.int16)).view(torch.bfloat16)
        hf = torch.from_numpy(np.array(H_test[t["offset"]:t["offset"] + n], copy=True).view(np.int16)).view(torch.bfloat16)
        lang = np.full(n, "", dtype=object)
        if "segments" in s:
            bounds = [x[0] for x in s["segments"]] + [n]
            for (st_, lg), en in zip(s["segments"], bounds[1:]):
                lang[st_:en] = lg
        toks = torch.tensor(s["token_ids"], device=dev)
        parts = defaultdict(list)
        for c in range(0, n, a.chunk):
            e = min(n, c + a.chunk)
            lb = torch.log_softmax(hb[c:e].to(dev, torch.float32) @ E.T, -1)
            lf = torch.log_softmax(hf[c:e].to(dev, torch.float32) @ E.T, -1)
            parts["kl"].append((lb.exp() * (lb - lf)).sum(-1).cpu())
            parts["rkl"].append((lf.exp() * (lf - lb)).sum(-1).cpu())
            parts["agree"].append((lb.argmax(-1) == lf.argmax(-1)).cpu())
            nxt = torch.arange(c + 1, e + 1, device=dev).clamp(max=n - 1)
            parts["our_lp"].append(lb.gather(1, toks[nxt].unsqueeze(1)).squeeze(1).cpu())  # logprob of token p+1
        our = torch.cat(parts["our_lp"]).numpy().astype(np.float64)
        our[-1] = np.nan
        vl = np.full(n, np.nan)
        vl[:-1] = lp_ref_vllm[r["offset"] + 1:r["offset"] + n]
        pos = np.arange(n)
        cols["kind"].append(np.full(n, s.get("kind", ""), dtype=object))
        cols["lang"].append(lang)
        cols["pos"].append(pos)
        cols["is_resp"].append(pos >= s.get("response_start", n + 1))
        cols["kl"].append(torch.cat(parts["kl"]).numpy().astype(np.float64))
        cols["rkl"].append(torch.cat(parts["rkl"]).numpy().astype(np.float64))
        cols["agree"].append(torch.cat(parts["agree"]).numpy())
        cols["our_lp"].append(our)
        cols["vl"].append(vl)
    kinds, langs, pos, is_resp = (np.concatenate(cols[k]) for k in ("kind", "lang", "pos", "is_resp"))
    klv, rkl, agree, our_lp, vl = (np.concatenate(cols[k]) for k in ("kl", "rkl", "agree", "our_lp", "vl"))
    out_dir = Path(a.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(out_dir / "per_position.npz", kind=kinds.astype(str), lang=langs.astype(str), pos=pos,
                        is_response=is_resp, kl=klv, rkl=rkl, top1_agree=agree)

    def stats(mask):
        m = mask & np.isfinite(klv)
        if not m.any():
            return None
        return {"positions": int(m.sum()), "kl_mean": float(klv[m].mean()), "kl_p50": float(np.median(klv[m])),
                "kl_p99": float(np.quantile(klv[m], 0.99)), "kl_max": float(klv[m].max()),
                "reverse_kl_mean": float(rkl[m].mean()), "top1_agreement": float(agree[m].mean())}
    every = np.ones_like(klv, dtype=bool)
    summary = {"ref": str(ref_dir), "test": str(test_dir), "threshold_nats_per_token": 0.02,
               "all": stats(every), "by_kind": {k: stats(kinds == k) for k in sorted(set(kinds))},
               "by_position_bucket": {f"{b * 512}-{b * 512 + 511}": stats((pos // 512) == b)
                                      for b in range(int(pos.max()) // 512 + 1)},
               "by_language": {l: stats(langs == l) for l in sorted(set(langs)) if l},
               "math_response_by_bucket": {f"{b * 512}-{b * 512 + 511}": stats((kinds == "math") & is_resp & ((pos // 512) == b))
                                           for b in range(int(pos.max()) // 512 + 1)}}
    ok = np.isfinite(our_lp) & np.isfinite(vl)
    summary["lm_head_reconstruction"] = {"positions": int(ok.sum()), "max_abs_diff_vs_vllm_logprob": float(np.abs(our_lp[ok] - vl[ok]).max()),
                                         "mean_abs_diff": float(np.abs(our_lp[ok] - vl[ok]).mean())}
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=1))
    print(json.dumps({k: summary[k] for k in ("all", "by_kind", "lm_head_reconstruction")}, indent=1))


# ---------------------------------------------------------------- scoring and analysis (CPU)


def score_gsm8k(a) -> None:
    gold = {r["id"]: r["gold"] for r in read_jsonl(DATA / "fp8check/gsm8k_250.jsonl")}
    runs = {}
    for spec in a.pred:
        name, path = spec.split("=", 1)
        runs[name] = {r["id"]: r for r in read_jsonl(path)}
    res = {}
    correct = {}
    for name, rows in runs.items():
        c = {i: same_number(extract_answer(r["text"]), gold[i]) for i, r in rows.items()}
        correct[name] = c
        res[name] = {"n": len(c), "accuracy": 100 * sum(c.values()) / len(c),
                     "truncated": sum(r["finish_reason"] == "length" for r in rows.values()),
                     "mean_out_tokens": float(np.mean([len(r["output_token_ids"]) for r in rows.values()]))}
    if len(runs) == 2:
        (na, ca), (nb, cb) = correct.items()
        ids = sorted(set(ca) & set(cb))
        res["paired"] = {"a": na, "b": nb, "diff_points": res[na]["accuracy"] - res[nb]["accuracy"],
                         "both_right": sum(ca[i] and cb[i] for i in ids), "only_a": sum(ca[i] and not cb[i] for i in ids),
                         "only_b": sum(cb[i] and not ca[i] for i in ids), "both_wrong": sum(not ca[i] and not cb[i] for i in ids),
                         "identical_text": sum(runs[na][i]["text"] == runs[nb][i]["text"] for i in ids),
                         "threshold_points": 1.5}
        res["paired"]["pass"] = abs(res["paired"]["diff_points"]) <= 1.5
    Path(a.out).write_text(json.dumps(res, indent=1))
    print(json.dumps(res, indent=1))


def analyze_trunc(a) -> None:
    rows = read_jsonl(a.generated)
    trunc = [r for r in rows if r["finish_reason"] == "length"]
    scores = [loop_score(r["output_token_ids"]) for r in trunc]
    loops = sum(s >= 0.5 for s in scores)
    by_lang = Counter(r.get("language", "") for r in trunc)
    tot_lang = Counter(r.get("language", "") for r in rows)
    out = {"n": len(rows), "truncated": len(trunc), "truncated_share": len(trunc) / max(1, len(rows)),
           "truncated_loops(score>=0.5)": loops, "truncated_non_loops": len(trunc) - loops,
           "loop_score_quartiles": [float(q) for q in np.quantile(scores, [0.25, 0.5, 0.75])] if scores else None,
           "truncated_by_language": {k: f"{by_lang[k]}/{tot_lang[k]}" for k in sorted(tot_lang)},
           "output_tokens_total": sum(len(r["output_token_ids"]) for r in rows),
           "output_tokens_in_truncated": sum(len(r["output_token_ids"]) for r in trunc)}
    examples = [{"id": r["id"], "language": r.get("language"), "loop_score": round(s, 3), "tail": r["text"][-300:]}
                for r, s in sorted(zip(trunc, scores), key=lambda x: x[1])][:: max(1, len(trunc) // 8)]
    Path(a.out).write_text(json.dumps({"summary": out, "examples": examples}, indent=1, ensure_ascii=False))
    print(json.dumps(out, indent=1))


def main() -> None:
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="op", required=True)
    s = sub.add_parser("prep-gsm8k"); s.add_argument("--seed", type=int, default=0); s.add_argument("--dataset-sha", default="740312add88f781978c0658806c59bc2815b9866")
    s = sub.add_parser("prep-math"); s.add_argument("--seed", type=int, default=0); s.add_argument("--n", type=int, default=200)
    s = sub.add_parser("prep-tf"); s.add_argument("--model", required=True); s.add_argument("--tokens", type=int, default=1_000_000)
    s.add_argument("--window", type=int, default=4000); s.add_argument("--seed", type=int, default=0)
    for name in ("generate", "capture"):
        s = sub.add_parser(name)
        s.add_argument("--engine", required=True)
        s.add_argument("--override", nargs="*", default=[], help="engine KEY=VAL (capacity knobs; diagnostic arms)")
        s.add_argument("--limit", type=int, default=0)
        s.add_argument("--out", required=True)
        if name == "generate":
            s.add_argument("--prompts", required=True)
            s.add_argument("--sampling", choices=["greedy", "spec"], required=True)
            s.add_argument("--max-tokens", type=int, required=True)
            s.add_argument("--sampling-config", default=str(PROD_CFG), help="config whose sampling block --sampling spec uses")
        else:
            s.add_argument("--sequences", nargs="+", required=True)
    s = sub.add_parser("seqs-from"); s.add_argument("--generated", required=True); s.add_argument("--out", required=True)
    s = sub.add_parser("kl"); s.add_argument("--ref", required=True); s.add_argument("--test", required=True)
    s.add_argument("--sequences", nargs="+", required=True); s.add_argument("--model", required=True)
    s.add_argument("--chunk", type=int, default=512); s.add_argument("--out", required=True); s.add_argument("--device", default="cuda")
    s = sub.add_parser("score-gsm8k"); s.add_argument("--pred", nargs="+", required=True, help="name=path"); s.add_argument("--out", required=True)
    s = sub.add_parser("analyze-trunc"); s.add_argument("--generated", required=True); s.add_argument("--out", required=True)
    a = p.parse_args()
    {"prep-gsm8k": prep_gsm8k, "prep-math": prep_math, "prep-tf": prep_tf, "generate": generate,
     "seqs-from": seqs_from, "capture": capture, "kl": kl, "score-gsm8k": score_gsm8k,
     "analyze-trunc": analyze_trunc}[a.op](a)


if __name__ == "__main__":
    main()