"""Exactly paired budget curve: cuts long generations (run with gen.py --save-tokens) to a smaller token budget.
Generation is autoregressive with a fixed seed, so the first B tokens of a 32k run are what a B-token run would
have produced; a record that finished within B tokens is unchanged, a longer one becomes a truncated record of B.
usage: budget_view.py --src results/m180_32k/mini.jsonl --budget 16384 --out results/m180_16k/mini.jsonl
"""
import argparse
import json
import os

from transformers import AutoTokenizer

ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--budget", type=int, required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
a = ap.parse_args()
tok = AutoTokenizer.from_pretrained(a.tokenizer)
os.makedirs(os.path.dirname(a.out), exist_ok=True)
cut = 0
with open(a.src, encoding="utf-8") as f, open(a.out, "w", encoding="utf-8") as out:
    for l in f:
        r = json.loads(l)
        if r["n_tokens"] > a.budget:
            ids = r["tokens"][: a.budget]
            r.update(completion=tok.decode(ids), tokens=ids, n_tokens=a.budget, truncated=True, stop_type="limit")
            cut += 1
        out.write(json.dumps(r, ensure_ascii=False) + "\n")
print(f"{a.src} -> budget {a.budget}: {cut} records cut")