LTA_openwebtext_dualt/scripts/apple_to_apple_lta_checks.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import json
import math
import sys
from collections import Counter
from pathlib import Path
from typing import Any, Iterable

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader


REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from eval import build_model_from_ckpt
from flowtext_lab.bridges import make_dirichlet_bridge_batch
from flowtext_lab.data import EosPadCollator, WrappedStreamingTextSequenceDataset, iter_text_records
from flowtext_lab.decode import sample_noise_simplex, state_for_model
from flowtext_lab.tokenization import BpeTextTokenizer
from train import TokenizedTextCollator, load_tokenized_hf_dataset


def token_piece(tok: BpeTextTokenizer, idx: int) -> str:
    raw = getattr(tok, "tokenizer", None)
    id_to_token = getattr(raw, "id_to_token", None)
    if callable(id_to_token):
        piece = id_to_token(int(idx))
        if piece is not None:
            return str(piece)
    return tok.decode([int(idx)], stop_at_eos=False, skip_special_tokens=False)


def token_text(tok: BpeTextTokenizer, idx: int) -> str:
    return tok.decode([int(idx)], stop_at_eos=False, skip_special_tokens=False)


def compact_piece(s: str) -> str:
    return s.replace("\n", "\\n").replace("\t", "\\t")


def load_batch(
    *,
    data_path: str,
    tokenizer: BpeTextTokenizer,
    max_len: int,
    batch_size: int,
    mode: str,
    text_column: str | None,
    openwebtext_split: str,
    wrap_mode: str,
    max_records: int,
    tokenized_pad_token: str,
) -> dict[str, torch.Tensor]:
    if mode == "tokenized_hf":
        ds = load_tokenized_hf_dataset(data_path, max_records=max_records)
        pad_id = tokenizer.pad_id if tokenized_pad_token == "pad" and tokenizer.pad_id is not None else tokenizer.eos_id
        collate = TokenizedTextCollator(pad_id, max_len=max_len)
        examples = [ds[i] for i in range(min(batch_size, len(ds)))]
        return collate(examples)
    if mode != "wrap":
        raise ValueError(f"unknown data mode: {mode}")
    ds = WrappedStreamingTextSequenceDataset(
        data_path,
        tokenizer,
        max_len=max_len,
        text_column=text_column,
        openwebtext_split=openwebtext_split,
        max_records_per_epoch=max_records,
        wrap_mode=wrap_mode,
    )
    loader = DataLoader(ds, batch_size=batch_size, collate_fn=EosPadCollator(tokenizer.eos_id, max_len=max_len))
    return next(iter(loader))


def iter_record_lengths(
    *,
    data_path: str,
    tokenizer: BpeTextTokenizer,
    mode: str,
    text_column: str | None,
    openwebtext_split: str,
    max_records: int,
) -> Iterable[int]:
    if mode == "tokenized_hf":
        ds = load_tokenized_hf_dataset(data_path, max_records=max_records)
        for ex in ds:
            raw = ex["input_ids"]
            if hasattr(raw, "tolist"):
                raw = raw.tolist()
            yield len(raw)
        return
    for i, text in enumerate(
        iter_text_records(
            data_path,
            text_column=text_column,
            openwebtext_split=openwebtext_split,
            detokenizer="auto",
        )
    ):
        if i >= max_records:
            break
        ids = tokenizer.encode(text, add_eos=False, add_special_tokens=False)
        yield len(ids)


def rate_summary(values: list[float]) -> dict[str, float]:
    if not values:
        return {"mean": 0.0, "min": 0.0, "p50": 0.0, "p90": 0.0, "p99": 0.0, "max": 0.0}
    vals = sorted(float(x) for x in values)
    n = len(vals)

    def q(p: float) -> float:
        return vals[min(n - 1, max(0, int(round((n - 1) * p))))]

    return {
        "mean": float(sum(vals) / n),
        "min": float(vals[0]),
        "p50": float(q(0.5)),
        "p90": float(q(0.9)),
        "p99": float(q(0.99)),
        "max": float(vals[-1]),
    }


def distribution_entropy_from_counts(counts: Counter[int]) -> float:
    total = sum(counts.values())
    if total <= 0:
        return 0.0
    out = 0.0
    for c in counts.values():
        p = c / total
        out -= p * math.log(max(p, 1e-12))
    return float(out)


def token_feature_rates(ids: torch.Tensor, tok: BpeTextTokenizer) -> dict[str, float]:
    flat = [int(x) for x in ids.reshape(-1).tolist()]
    if not flat:
        return {}
    pieces = [token_piece(tok, x) for x in flat]
    texts = [token_text(tok, x) for x in flat]
    specials = {tok.eos_id, tok.bos_id, tok.unk_id}
    if tok.pad_id is not None:
        specials.add(tok.pad_id)
    denom = len(flat)
    normal = [i for i, x in enumerate(flat) if x not in specials]
    normal_denom = max(len(normal), 1)
    return {
        "bert_hash_rate": sum(pieces[i].startswith("##") for i in normal) / normal_denom,
        "spm_cont_rate": sum((not pieces[i].startswith("▁")) and (not pieces[i].startswith("<")) for i in normal) / normal_denom,
        "single_char_rate": sum(len(texts[i].strip()) == 1 for i in normal) / normal_denom,
        "digit_piece_rate": sum(any(ch.isdigit() for ch in pieces[i]) for i in normal) / normal_denom,
        "url_piece_rate": sum(("http" in pieces[i].lower() or "www" in pieces[i].lower() or ".com" in pieces[i].lower()) for i in normal) / normal_denom,
        "special_rate": sum(x in specials for x in flat) / denom,
    }


def command_data(args: argparse.Namespace) -> None:
    tok = BpeTextTokenizer.from_file(args.tokenizer_path)
    batch = load_batch(
        data_path=args.data_path,
        tokenizer=tok,
        max_len=args.max_len,
        batch_size=args.n_sequences,
        mode=args.data_mode,
        text_column=args.text_column,
        openwebtext_split=args.openwebtext_split,
        wrap_mode=args.wrap_mode,
        max_records=args.max_records,
        tokenized_pad_token=args.tokenized_pad_token,
    )
    ids = batch["ids"]
    attn = batch.get("attn_mask", torch.ones_like(ids, dtype=torch.bool))
    valid_ids = ids[attn]
    counts = Counter(int(x) for x in valid_ids.tolist())
    top = [
        {
            "id": int(i),
            "piece": compact_piece(token_piece(tok, int(i))),
            "text": compact_piece(token_text(tok, int(i))),
            "count": int(c),
            "rate": float(c / max(valid_ids.numel(), 1)),
        }
        for i, c in counts.most_common(args.top_k)
    ]
    seq_lens = attn.long().sum(dim=1).tolist()
    internal = ids[:, 1:-1] if ids.size(1) > 2 else ids[:, :0]
    internal_attn = attn[:, 1:-1] if attn.size(1) > 2 else attn[:, :0]
    eos_internal = ((internal == tok.eos_id) & internal_attn).long().sum(dim=1).tolist()
    pad_internal = []
    if tok.pad_id is not None:
        pad_internal = ((internal == tok.pad_id) & internal_attn).long().sum(dim=1).tolist()
    pos0 = Counter(int(x) for x in ids[:, 0].tolist())
    last_valid = []
    for row, mask in zip(ids, attn):
        idx = int(mask.long().sum().item()) - 1
        if idx >= 0:
            last_valid.append(int(row[idx].item()))
    last_counts = Counter(last_valid)
    record_lengths = list(
        iter_record_lengths(
            data_path=args.data_path,
            tokenizer=tok,
            mode=args.data_mode,
            text_column=args.text_column,
            openwebtext_split=args.openwebtext_split,
            max_records=args.max_records,
        )
    )
    out = {
        "name": args.name,
        "data_path": args.data_path,
        "data_mode": args.data_mode,
        "tokenizer_path": args.tokenizer_path,
        "vocab_size": tok.vocab_size,
        "bos_id": tok.bos_id,
        "bos_piece": token_piece(tok, tok.bos_id),
        "eos_id": tok.eos_id,
        "eos_piece": token_piece(tok, tok.eos_id),
        "pad_id": tok.pad_id,
        "n_sequences": int(ids.size(0)),
        "max_len": args.max_len,
        "sequence_len": rate_summary([float(x) for x in seq_lens]),
        "record_token_len_no_special_no_eos": rate_summary([float(x) for x in record_lengths]),
        "internal_eos_per_seq": rate_summary([float(x) for x in eos_internal]),
        "internal_pad_per_seq": rate_summary([float(x) for x in pad_internal]) if pad_internal else None,
        "pos0_top": [
            {"id": i, "piece": compact_piece(token_piece(tok, i)), "count": c, "rate": c / max(ids.size(0), 1)}
            for i, c in pos0.most_common(args.top_k)
        ],
        "last_valid_top": [
            {"id": i, "piece": compact_piece(token_piece(tok, i)), "count": c, "rate": c / max(len(last_valid), 1)}
            for i, c in last_counts.most_common(args.top_k)
        ],
        "unigram_entropy": distribution_entropy_from_counts(counts),
        "token_feature_rates": token_feature_rates(valid_ids, tok),
        "top_unigram": top,
    }
    Path(args.out_json).parent.mkdir(parents=True, exist_ok=True)
    Path(args.out_json).write_text(json.dumps(out, indent=2, ensure_ascii=False), encoding="utf-8")
    print(json.dumps(out, indent=2, ensure_ascii=False), flush=True)


def ckpt_arg(ckpt_args: dict[str, Any], key: str, default: Any) -> Any:
    return ckpt_args.get(key, default)


def make_bridge_for_eval(
    *,
    ids: torch.Tensor,
    attn: torch.Tensor,
    ckpt_args: dict[str, Any],
    vocab_size: int,
    t_value: float,
    force_mask_ratio: float | None,
    eps: float,
) -> Any:
    return make_dirichlet_bridge_batch(
        ids=ids,
        attn_mask=attn,
        vocab_size=vocab_size,
        target_prob=float(ckpt_arg(ckpt_args, "target_prob", 1.0)),
        min_t=float(ckpt_arg(ckpt_args, "min_t", 0.0)),
        max_t=float(ckpt_arg(ckpt_args, "max_t", 1.0)),
        min_mask_ratio=float(ckpt_arg(ckpt_args, "min_mask_ratio", 0.1)),
        max_mask_ratio=float(ckpt_arg(ckpt_args, "max_mask_ratio", 1.0)),
        wrong_token_replace_prob=ckpt_arg(ckpt_args, "wrong_token_replace_prob", "0.0"),
        wrong_token_schedule=str(ckpt_arg(ckpt_args, "wrong_token_schedule", "constant")),
        wrong_token_exp_k=float(ckpt_arg(ckpt_args, "wrong_token_exp_k", 1.0)),
        dirichlet_concentration_min=float(ckpt_arg(ckpt_args, "dirichlet_concentration_min", 1.0)),
        dirichlet_concentration_max=float(ckpt_arg(ckpt_args, "dirichlet_concentration_max", 1024.0)),
        eps=eps,
        state_format=str(ckpt_arg(ckpt_args, "state_format", ckpt_arg(ckpt_args, "input_format", "prob"))),
        dirichlet_endpoint_mode=str(ckpt_arg(ckpt_args, "dirichlet_endpoint_mode", "bernoulli_wrong")),
        dirichlet_semantic_t_mode=str(ckpt_arg(ckpt_args, "dirichlet_semantic_t_mode", "same")),
        dirichlet_semantic_t_value=float(ckpt_arg(ckpt_args, "dirichlet_semantic_t_value", 0.0)),
        dirichlet_semantic_t_curve=str(ckpt_arg(ckpt_args, "dirichlet_semantic_t_curve", "linear")),
        dirichlet_semantic_t_power=float(ckpt_arg(ckpt_args, "dirichlet_semantic_t_power", 1.0)),
        dirichlet_support_t_curve=str(ckpt_arg(ckpt_args, "dirichlet_support_t_curve", "linear")),
        dirichlet_support_t_power=float(ckpt_arg(ckpt_args, "dirichlet_support_t_power", 1.0)),
        endpoint_sequence_random_prob_alpha=float(ckpt_arg(ckpt_args, "endpoint_sequence_random_prob_alpha", 0.0)),
        categorical_wrong_from_full_vocab=bool(ckpt_arg(ckpt_args, "categorical_wrong_from_full_vocab", False)),
        categorical_wrong_from_batch_valid_tokens=bool(ckpt_arg(ckpt_args, "categorical_wrong_from_batch_valid_tokens", False)),
        categorical_wrong_basin_token_ids=ckpt_arg(ckpt_args, "categorical_wrong_basin_token_ids", ""),
        categorical_wrong_basin_prob=float(ckpt_arg(ckpt_args, "categorical_wrong_basin_prob", 0.0)),
        categorical_wrong_unigram_prob=float(ckpt_arg(ckpt_args, "categorical_wrong_unigram_prob", 0.0)),
        categorical_wrong_uniform_prob=float(ckpt_arg(ckpt_args, "categorical_wrong_uniform_prob", 0.0)),
        categorical_wrong_prob_floor=float(ckpt_arg(ckpt_args, "categorical_wrong_prob_floor", 0.0)),
        categorical_gold_prob_floor=float(ckpt_arg(ckpt_args, "categorical_gold_prob_floor", 0.0)),
        categorical_gold_prob_ceil=float(ckpt_arg(ckpt_args, "categorical_gold_prob_ceil", 1.0)),
        simplex_bridge_sampler=str(ckpt_arg(ckpt_args, "simplex_bridge_sampler", "dirichlet")),
        logistic_normal_sigma_min=float(ckpt_arg(ckpt_args, "logistic_normal_sigma_min", 0.18)),
        logistic_normal_sigma_max=float(ckpt_arg(ckpt_args, "logistic_normal_sigma_max", 2.2)),
        logistic_normal_tau_min=float(ckpt_arg(ckpt_args, "logistic_normal_tau_min", 0.65)),
        logistic_normal_tau_max=float(ckpt_arg(ckpt_args, "logistic_normal_tau_max", 1.15)),
        force_t=t_value,
        force_mask_ratio=force_mask_ratio,
        mask_ratio_floor_schedule=str(ckpt_arg(ckpt_args, "mask_ratio_floor_schedule", "none")),
        mask_mixture_original_prob=float(ckpt_arg(ckpt_args, "mask_mixture_original_prob", 0.0)),
        mask_mixture_lowk_prob=float(ckpt_arg(ckpt_args, "mask_mixture_lowk_prob", 0.0)),
        mask_mixture_lowcorrupt_prob=float(ckpt_arg(ckpt_args, "mask_mixture_lowcorrupt_prob", 0.0)),
        mask_mixture_block_prob=float(ckpt_arg(ckpt_args, "mask_mixture_block_prob", 0.0)),
        mask_mixture_all_prob=float(ckpt_arg(ckpt_args, "mask_mixture_all_prob", 0.0)),
        mask_mixture_lowk_clean_tokens=ckpt_arg(ckpt_args, "mask_mixture_lowk_clean_tokens", "1,2,4,8,16,32,64"),
        mask_mixture_lowcorrupt_tokens=ckpt_arg(ckpt_args, "mask_mixture_lowcorrupt_tokens", "1,2,4,8,16,32,64"),
        mask_mixture_block_tokens=ckpt_arg(ckpt_args, "mask_mixture_block_tokens", "64,128"),
        clean_state_mode=str(ckpt_arg(ckpt_args, "clean_state_mode", "onehot")),
        return_dense_targets=False,
    )


def masked_loss_acc(logits: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> dict[str, float]:
    flat_mask = mask.reshape(-1)
    if not bool(flat_mask.any().item()):
        return {"nll": 0.0, "ppl": 1.0, "acc": 0.0, "tokens": 0}
    flat_logits = logits.reshape(-1, logits.size(-1))[flat_mask]
    flat_target = target.reshape(-1)[flat_mask]
    loss = F.cross_entropy(flat_logits, flat_target, reduction="mean")
    pred = flat_logits.argmax(dim=-1)
    acc = (pred == flat_target).float().mean()
    return {
        "nll": float(loss.detach().cpu()),
        "ppl": float(torch.exp(loss.clamp(max=50)).detach().cpu()),
        "acc": float(acc.detach().cpu()),
        "tokens": int(flat_mask.sum().detach().cpu()),
    }


@torch.inference_mode()
def command_teacher(args: argparse.Namespace) -> None:
    tok = BpeTextTokenizer.from_file(args.tokenizer_path)
    device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu else "cpu")
    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    ckpt_args = dict(ckpt.get("args", {}))
    model = build_model_from_ckpt(ckpt, tok.vocab_size, args.max_len, device).eval()
    batch = load_batch(
        data_path=args.data_path,
        tokenizer=tok,
        max_len=args.max_len,
        batch_size=args.batch_size,
        mode=args.data_mode,
        text_column=args.text_column,
        openwebtext_split=args.openwebtext_split,
        wrap_mode=args.wrap_mode,
        max_records=args.max_records,
        tokenized_pad_token=args.tokenized_pad_token,
    )
    ids = batch["ids"].to(device)
    attn = batch.get("attn_mask", torch.ones_like(ids, dtype=torch.bool)).to(device)
    rows = []
    for t_value in [float(x) for x in args.t_values.split(",") if x.strip()]:
        torch.manual_seed(args.seed + int(round(t_value * 1000000)))
        bridge = make_bridge_for_eval(
            ids=ids,
            attn=attn,
            ckpt_args=ckpt_args,
            vocab_size=tok.vocab_size,
            t_value=t_value,
            force_mask_ratio=args.force_mask_ratio,
            eps=args.eps,
        )
        model_t = bridge.t
        logits = model(state_for_model(model, bridge.state, args.eps), model_t, attn).float()
        valid = attn
        corrupt = bridge.corrupt_mask & attn
        pos0_pred = logits[:, 0].argmax(dim=-1)
        last_pred = []
        for b in range(ids.size(0)):
            last = int(attn[b].long().sum().item()) - 1
            last_pred.append(int(logits[b, last].argmax().detach().cpu()) if last >= 0 else -1)
        pos0_counts = Counter(int(x) for x in pos0_pred.detach().cpu().tolist())
        last_counts = Counter(last_pred)
        probs = F.softmax(logits, dim=-1)
        rows.append(
            {
                "name": args.name,
                "checkpoint": args.checkpoint,
                "ckpt_step": int(ckpt.get("step", -1)),
                "t": t_value,
                "force_mask_ratio": args.force_mask_ratio,
                "corrupt_frac": float(corrupt.float().mean().detach().cpu()),
                "wrong_frac": float((bridge.wrong_mask & attn).float().sum().detach().cpu() / attn.float().sum().clamp_min(1).detach().cpu()),
                "valid": masked_loss_acc(logits, ids, valid),
                "corrupt": masked_loss_acc(logits, ids, corrupt),
                "dist_entropy": float((-(probs.clamp_min(args.eps) * probs.clamp_min(args.eps).log()).sum(dim=-1)[valid]).mean().detach().cpu()),
                "mean_maxp": float(probs.max(dim=-1).values[valid].mean().detach().cpu()),
                "pos0_gold_id": int(ids[0, 0].detach().cpu()),
                "pos0_gold_piece": token_piece(tok, int(ids[0, 0].detach().cpu())),
                "pos0_top": [
                    {"id": i, "piece": compact_piece(token_piece(tok, i)), "count": c, "rate": c / max(ids.size(0), 1)}
                    for i, c in pos0_counts.most_common(5)
                ],
                "last_top": [
                    {"id": i, "piece": compact_piece(token_piece(tok, i)), "count": c, "rate": c / max(ids.size(0), 1)}
                    for i, c in last_counts.most_common(5)
                ],
            }
        )
    out = Path(args.out_json)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(rows, indent=2, ensure_ascii=False), encoding="utf-8")
    with out.with_suffix(".tsv").open("w", newline="", encoding="utf-8") as f:
        fields = [
            "name",
            "ckpt_step",
            "t",
            "force_mask_ratio",
            "corrupt_frac",
            "wrong_frac",
            "valid_nll",
            "valid_acc",
            "corrupt_nll",
            "corrupt_acc",
            "dist_entropy",
            "mean_maxp",
            "pos0_gold_piece",
            "pos0_top",
            "last_top",
        ]
        writer = csv.DictWriter(f, fieldnames=fields, delimiter="\t")
        writer.writeheader()
        for row in rows:
            writer.writerow(
                {
                    "name": row["name"],
                    "ckpt_step": row["ckpt_step"],
                    "t": row["t"],
                    "force_mask_ratio": row["force_mask_ratio"],
                    "corrupt_frac": row["corrupt_frac"],
                    "wrong_frac": row["wrong_frac"],
                    "valid_nll": row["valid"]["nll"],
                    "valid_acc": row["valid"]["acc"],
                    "corrupt_nll": row["corrupt"]["nll"],
                    "corrupt_acc": row["corrupt"]["acc"],
                    "dist_entropy": row["dist_entropy"],
                    "mean_maxp": row["mean_maxp"],
                    "pos0_gold_piece": row["pos0_gold_piece"],
                    "pos0_top": " | ".join(f"{x['piece']}:{x['rate']:.2f}" for x in row["pos0_top"]),
                    "last_top": " | ".join(f"{x['piece']}:{x['rate']:.2f}" for x in row["last_top"]),
                }
            )
    for row in rows:
        print(
            f"{row['name']} step={row['ckpt_step']} t={row['t']:.4f} "
            f"valid_nll={row['valid']['nll']:.3f} valid_acc={row['valid']['acc']:.3f} "
            f"corrupt_nll={row['corrupt']['nll']:.3f} corrupt_acc={row['corrupt']['acc']:.3f} "
            f"pos0={row['pos0_top'][0]['piece']}:{row['pos0_top'][0]['rate']:.2f}",
            flush=True,
        )


def filter_top_p(probs: torch.Tensor, top_p: float, eps: float) -> torch.Tensor:
    if top_p >= 1.0:
        return probs
    sorted_vals, sorted_idx = torch.sort(probs, dim=-1, descending=True)
    total = sorted_vals.sum(dim=-1, keepdim=True).clamp_min(eps)
    remove = sorted_vals.cumsum(dim=-1) > top_p * total
    remove[..., 0] = False
    sorted_vals = sorted_vals.masked_fill(remove, 0.0)
    out = torch.zeros_like(probs).scatter(-1, sorted_idx, sorted_vals)
    return out / out.sum(dim=-1, keepdim=True).clamp_min(eps)


def distribution_metrics(probs: torch.Tensor, ids: torch.Tensor, tok: BpeTextTokenizer, prefix: str) -> dict[str, Any]:
    p = probs.clamp_min(1e-12)
    ent = float((-(p * p.log()).sum(dim=-1)).mean().detach().cpu())
    maxp, arg = probs.max(dim=-1)
    counts = Counter(int(x) for x in arg.reshape(-1).detach().cpu().tolist())
    return {
        f"{prefix}_entropy": ent,
        f"{prefix}_mean_top_mass": float(maxp.mean().detach().cpu()),
        f"{prefix}_argmax_token_entropy": distribution_entropy_from_counts(counts),
        f"{prefix}_argmax_top": [
            {"id": i, "piece": compact_piece(token_piece(tok, i)), "count": c, "rate": c / max(arg.numel(), 1)}
            for i, c in counts.most_common(8)
        ],
        **{f"{prefix}_{k}": v for k, v in token_feature_rates(arg.detach().cpu(), tok).items()},
    }


@torch.inference_mode()
def command_trace(args: argparse.Namespace) -> None:
    tok = BpeTextTokenizer.from_file(args.tokenizer_path)
    device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu else "cpu")
    torch.manual_seed(args.seed)
    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    model = build_model_from_ckpt(ckpt, tok.vocab_size, args.max_len, device).eval()
    eps = args.eps
    bs = args.batch_size
    probs = sample_noise_simplex(
        (bs, args.max_len),
        tok.vocab_size,
        device,
        eps,
        noise_mode="dirichlet",
        target_prob=1.0,
        noise_sigma=-1.0,
        dirichlet_concentration=args.concentration_min,
    )
    attn = torch.ones((bs, args.max_len), dtype=torch.bool, device=device)
    log_cmin = math.log(args.concentration_min)
    log_cmax = math.log(args.concentration_max)
    out = Path(args.out_jsonl)
    out.parent.mkdir(parents=True, exist_ok=True)
    snapshot = set(int(x) for x in args.trace_steps.split(",") if x.strip())
    last_endpoint = probs
    with out.open("w", encoding="utf-8") as f:
        for step in range(args.steps):
            support_t = (step + 1) / max(args.steps, 1)
            t = torch.full((bs,), support_t, dtype=torch.float32, device=device)
            logits = model(state_for_model(model, probs, eps), t, attn).float()
            endpoint = F.softmax(logits / args.endpoint_temp, dim=-1)
            endpoint = filter_top_p(endpoint, args.endpoint_top_p, eps)
            tau = args.gumbel_tau_start + support_t * (args.gumbel_tau_end - args.gumbel_tau_start)
            uniform = torch.rand_like(endpoint).clamp_(eps, 1.0 - eps)
            gumbel = -torch.log(-torch.log(uniform))
            projected = F.softmax((endpoint.clamp_min(eps).log() + gumbel) / max(tau, eps), dim=-1)
            last_endpoint = projected
            mean = (1.0 - support_t) / tok.vocab_size + support_t * projected
            mean = mean / mean.sum(dim=-1, keepdim=True).clamp_min(eps)
            conc = math.exp(log_cmin + support_t * (log_cmax - log_cmin))
            alpha = (mean * conc).clamp_min(eps)
            probs = torch._standard_gamma(alpha).clamp_min(eps)
            probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(eps)
            step_num = step + 1
            if step_num in snapshot or step_num == args.steps:
                row = {
                    "name": args.name,
                    "ckpt_step": int(ckpt.get("step", -1)),
                    "step": step_num,
                    "support_t": support_t,
                    "tau": tau,
                    "concentration": conc,
                }
                row.update(distribution_metrics(endpoint, endpoint.argmax(dim=-1), tok, "a"))
                row.update(distribution_metrics(projected, projected.argmax(dim=-1), tok, "e"))
                row.update(distribution_metrics(probs, probs.argmax(dim=-1), tok, "p"))
                for pos in [0, 1, args.max_len - 2, args.max_len - 1]:
                    a_id = int(endpoint[0, pos].argmax().detach().cpu())
                    e_id = int(projected[0, pos].argmax().detach().cpu())
                    p_id = int(probs[0, pos].argmax().detach().cpu())
                    row[f"pos{pos}_a"] = {"id": a_id, "piece": compact_piece(token_piece(tok, a_id)), "prob": float(endpoint[0, pos, a_id].detach().cpu())}
                    row[f"pos{pos}_e"] = {"id": e_id, "piece": compact_piece(token_piece(tok, e_id)), "prob": float(projected[0, pos, e_id].detach().cpu())}
                    row[f"pos{pos}_p"] = {"id": p_id, "piece": compact_piece(token_piece(tok, p_id)), "prob": float(probs[0, pos, p_id].detach().cpu())}
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
                print(
                    f"{args.name} step={step_num} aH={row['a_entropy']:.2f} eH={row['e_entropy']:.2f} pH={row['p_entropy']:.2f} "
                    f"a_top={row['a_argmax_top'][0]['piece']}:{row['a_argmax_top'][0]['rate']:.2f} "
                    f"p_top={row['p_argmax_top'][0]['piece']}:{row['p_argmax_top'][0]['rate']:.2f}",
                    flush=True,
                )
    if args.final_out:
        final_probs = 0.5 * probs + 0.5 * last_endpoint
        ids = final_probs.argmax(dim=-1).detach().cpu().tolist()
        Path(args.final_out).write_text("\n\n".join(tok.decode(row, stop_at_eos=False, skip_special_tokens=False) for row in ids), encoding="utf-8")


def main() -> None:
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    data = sub.add_parser("data")
    data.add_argument("--name", required=True)
    data.add_argument("--data_path", required=True)
    data.add_argument("--tokenizer_path", required=True)
    data.add_argument("--out_json", required=True)
    data.add_argument("--data_mode", choices=["wrap", "tokenized_hf"], default="wrap")
    data.add_argument("--text_column", default=None)
    data.add_argument("--openwebtext_split", default="all")
    data.add_argument("--wrap_mode", default="stream")
    data.add_argument("--tokenized_pad_token", default="pad")
    data.add_argument("--max_len", type=int, default=1024)
    data.add_argument("--n_sequences", type=int, default=2048)
    data.add_argument("--max_records", type=int, default=20000)
    data.add_argument("--top_k", type=int, default=24)
    data.set_defaults(func=command_data)

    teacher = sub.add_parser("teacher")
    teacher.add_argument("--name", required=True)
    teacher.add_argument("--checkpoint", required=True)
    teacher.add_argument("--data_path", required=True)
    teacher.add_argument("--tokenizer_path", required=True)
    teacher.add_argument("--out_json", required=True)
    teacher.add_argument("--data_mode", choices=["wrap", "tokenized_hf"], default="wrap")
    teacher.add_argument("--text_column", default=None)
    teacher.add_argument("--openwebtext_split", default="all")
    teacher.add_argument("--wrap_mode", default="stream")
    teacher.add_argument("--tokenized_pad_token", default="pad")
    teacher.add_argument("--max_len", type=int, default=1024)
    teacher.add_argument("--batch_size", type=int, default=8)
    teacher.add_argument("--max_records", type=int, default=20000)
    teacher.add_argument("--t_values", default="0.0,0.0078125,0.03125,0.125,0.5,1.0")
    teacher.add_argument("--force_mask_ratio", type=float, default=None)
    teacher.add_argument("--seed", type=int, default=20260525)
    teacher.add_argument("--eps", type=float, default=1e-8)
    teacher.add_argument("--cpu", action="store_true")
    teacher.set_defaults(func=command_teacher)

    trace = sub.add_parser("trace")
    trace.add_argument("--name", required=True)
    trace.add_argument("--checkpoint", required=True)
    trace.add_argument("--tokenizer_path", required=True)
    trace.add_argument("--out_jsonl", required=True)
    trace.add_argument("--final_out", default="")
    trace.add_argument("--max_len", type=int, default=1024)
    trace.add_argument("--batch_size", type=int, default=2)
    trace.add_argument("--steps", type=int, default=128)
    trace.add_argument("--trace_steps", default="1,2,4,8,16,32,64,96,128")
    trace.add_argument("--concentration_min", type=float, default=30522)
    trace.add_argument("--concentration_max", type=float, default=61044)
    trace.add_argument("--endpoint_temp", type=float, default=1.45)
    trace.add_argument("--endpoint_top_p", type=float, default=0.95)
    trace.add_argument("--gumbel_tau_start", type=float, default=1.0)
    trace.add_argument("--gumbel_tau_end", type=float, default=0.2)
    trace.add_argument("--seed", type=int, default=20260525)
    trace.add_argument("--eps", type=float, default=1e-8)
    trace.add_argument("--cpu", action="store_true")
    trace.set_defaults(func=command_trace)

    args = ap.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()