LTA_openwebtext_dualt/scripts/eval_c1024_decode_sweep_20260507.py

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

import argparse
import csv
import json
import math
import re
import sys
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer


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

from flowtext_lab.decode import sample_noise_simplex, state_for_model
from flowtext_lab.genppl import summarize_token_diversity
from flowtext_lab.model import EndpointPredictor
from flowtext_lab.tokenization import BpeTextTokenizer


SPACE_RE = re.compile(r"\s+")


@dataclass(frozen=True)
class DecodeConfig:
    name: str
    model_t_mode: str
    support_power: float
    semantic_power: float
    final_from: str
    endpoint_temp: float = 1.3
    concentration_min: float = 1.0
    concentration_max: float = 1024.0
    update: str = "resample"


def default_configs() -> list[DecodeConfig]:
    return [
        DecodeConfig("baseline_const05_sem15_blend_c1024_t1p3", "const05", 1.0, 1.5, "blend"),
        DecodeConfig("match_post_sem1_blend_c1024_t1p3", "post", 1.0, 1.0, "blend"),
        DecodeConfig("match_post_sem1_state_c1024_t1p3", "post", 1.0, 1.0, "state"),
        DecodeConfig("match_post_sem1_endpoint_c1024_t1p3", "post", 1.0, 1.0, "endpoint"),
        DecodeConfig("pre_sem1_blend_c1024_t1p3", "pre", 1.0, 1.0, "blend"),
        DecodeConfig("const1_sem1_blend_c1024_t1p3", "const1", 1.0, 1.0, "blend"),
        DecodeConfig("const05_sem1_blend_c1024_t1p3", "const05", 1.0, 1.0, "blend"),
        DecodeConfig("match_post_sem1_blend_c1024_t1p8", "post", 1.0, 1.0, "blend", endpoint_temp=1.8),
        DecodeConfig("match_post_sem1_state_c1024_t1p8", "post", 1.0, 1.0, "state", endpoint_temp=1.8),
        DecodeConfig("match_post_sem1_blend_c256_t1p3", "post", 1.0, 1.0, "blend", concentration_max=256.0),
        DecodeConfig("match_post_sem1_blend_c64_t1p3", "post", 1.0, 1.0, "blend", concentration_max=64.0),
        DecodeConfig("match_post_sem1_blend_c16_t1p3", "post", 1.0, 1.0, "blend", concentration_max=16.0),
        DecodeConfig("slow_support_post_sem1_blend_c1024_t1p3", "post", 2.0, 1.0, "blend"),
        DecodeConfig("fast_support_post_sem1_blend_c1024_t1p3", "post", 0.5, 1.0, "blend"),
        DecodeConfig("match_post_sem1_mean_c1024_t1p3", "post", 1.0, 1.0, "blend", update="mean"),
    ]


def lm1b_detokenizer(text: str) -> str:
    text = text.replace("http : / / ", "http://")
    text = text.replace("https : / / ", "https://")
    text = re.sub(r" ' (s|m|re|ve|ll|d|t)\b", r"'\1", text)
    text = re.sub(r" '(s|m|re|ve|ll|d|t)\b", r"'\1", text)
    text = re.sub(r"n ' t\b", "n't", text)
    text = re.sub(r"n 't\b", "n't", text)
    text = re.sub(r"(\w+) ' (\w+)", r"\1' \2", text)
    text = re.sub(r" (\w+) \. ", r" \1. ", text)
    text = re.sub(r" (\w+) \.$", r" \1.", text)
    text = text.replace(" ? ", "? ")
    text = re.sub(r" \?$", "?", text)
    text = text.replace(" ! ", "! ")
    text = re.sub(r" \!$", "!", text)
    text = text.replace(" , ", ", ")
    text = text.replace(" : ", ": ")
    text = text.replace(" ; ", "; ")
    text = text.replace(" / ", "/")
    text = re.sub(r'" ([^"]+) "', r'"\1"', text)
    text = re.sub(r"' ([^']+) '", r"'\1'", text)
    text = re.sub(r"\( ([^\(\)]+) \)", r"(\1)", text)
    text = re.sub(r"\[ ([^\[\]]+) \]", r"[\1]", text)
    text = text.replace("$ ", "$")
    text = text.replace("£ ", "£")
    return SPACE_RE.sub(" ", text).strip()


def model_time(mode: str, step: int, steps: int, batch: int, device: torch.device) -> torch.Tensor:
    if mode == "pre":
        value = step / max(steps, 1)
    elif mode == "post":
        value = (step + 1) / max(steps, 1)
    elif mode == "linear":
        value = step / max(steps - 1, 1)
    elif mode == "const05":
        value = 0.5
    elif mode == "const1":
        value = 1.0
    elif mode == "const0":
        value = 0.0
    else:
        raise ValueError(f"unknown model_t_mode: {mode}")
    return torch.full((batch,), float(value), dtype=torch.float32, device=device)


def build_model(ckpt: dict, tokenizer: BpeTextTokenizer, device: torch.device) -> EndpointPredictor:
    args = ckpt.get("args", {})
    train_vocab = int(ckpt.get("train_vocab_size") or ckpt.get("vocab_size") or tokenizer.vocab_size)
    model = EndpointPredictor(
        vocab_size=train_vocab,
        max_len=int(args.get("max_len", 128)),
        d_model=int(args.get("d_model", 768)),
        n_heads=int(args.get("n_heads", 12)),
        n_layers=int(args.get("n_layers", 12)),
        dim_ff=int(args.get("dim_ff", 3072)),
        dropout=0.0,
        model_type=str(args.get("model_type", "ddit")),
        cond_dim=int(args.get("cond_dim", 128)),
        input_format=str(args.get("state_format", args.get("input_format", "prob"))),
    ).to(device)
    model.load_state_dict(ckpt["model"], strict=True)
    model.eval()
    return model


@torch.no_grad()
def decode_config(
    model: EndpointPredictor,
    tokenizer: BpeTextTokenizer,
    cfg: DecodeConfig,
    *,
    n_samples: int,
    batch_size: int,
    max_len: int,
    steps: int,
    seed: int,
    device: torch.device,
) -> tuple[list[list[int]], list[str], dict]:
    torch.manual_seed(seed)
    eps = 1e-8
    all_ids: list[list[int]] = []
    all_texts: list[str] = []
    remaining = n_samples
    while remaining > 0:
        bs = min(batch_size, remaining)
        probs = sample_noise_simplex(
            (bs, max_len),
            tokenizer.vocab_size,
            device,
            eps,
            noise_mode="dirichlet",
            target_prob=1.0,
            noise_sigma=-1.0,
            dirichlet_concentration=1.0,
        )
        attn = torch.ones((bs, max_len), dtype=torch.bool, device=device)
        last_endpoint = probs
        for step in range(steps):
            support_t = ((step + 1) / max(steps, 1)) ** cfg.support_power
            semantic_t = ((step + 1) / max(steps, 1)) ** cfg.semantic_power
            t = model_time(cfg.model_t_mode, step, steps, bs, device)
            logits = model(state_for_model(model, probs, eps), t, attn).float() / cfg.endpoint_temp
            endpoint = F.softmax(logits, dim=-1)
            last_endpoint = endpoint
            anchor = probs.clamp_min(eps)
            anchor = anchor / anchor.sum(dim=-1, keepdim=True).clamp_min(eps)
            forward_endpoint = (1.0 - semantic_t) * anchor + semantic_t * endpoint
            forward_endpoint = forward_endpoint.clamp_min(eps)
            forward_endpoint = forward_endpoint / forward_endpoint.sum(dim=-1, keepdim=True).clamp_min(eps)
            mean = (1.0 - support_t) / float(tokenizer.vocab_size) + support_t * forward_endpoint
            mean = mean.clamp_min(eps)
            mean = mean / mean.sum(dim=-1, keepdim=True).clamp_min(eps)
            if cfg.update == "mean":
                probs = mean
            elif cfg.update == "resample":
                log_min = math.log(max(cfg.concentration_min, eps))
                log_max = math.log(max(cfg.concentration_max, cfg.concentration_min))
                conc = math.exp(log_min + support_t * (log_max - log_min))
                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)
            else:
                raise ValueError(cfg.update)
        if cfg.final_from == "endpoint":
            final = last_endpoint
        elif cfg.final_from == "state":
            final = probs
        elif cfg.final_from == "blend":
            final = 0.5 * probs + 0.5 * last_endpoint
        else:
            raise ValueError(cfg.final_from)
        final = final / final.sum(dim=-1, keepdim=True).clamp_min(eps)
        ids = final.argmax(dim=-1).detach().cpu().tolist()
        all_ids.extend(ids)
        all_texts.extend(tokenizer.decode(row, stop_at_eos=False, skip_special_tokens=False) for row in ids)
        remaining -= bs
    decode = {
        "kind": "categorical_fullvocab",
        "config": cfg.name,
        "steps": steps,
        "model_t_mode": cfg.model_t_mode,
        "support_power": cfg.support_power,
        "semantic_power": cfg.semantic_power,
        "endpoint_temp": cfg.endpoint_temp,
        "concentration_min": cfg.concentration_min,
        "concentration_max": cfg.concentration_max,
        "final_from": cfg.final_from,
        "update": cfg.update,
        "n_samples": n_samples,
        "seed": seed,
    }
    return all_ids, all_texts, decode


@torch.no_grad()
def score_texts(texts: list[str], scorer, scorer_tok, *, batch_size: int, max_length: int, device: torch.device) -> dict:
    total_nll = 0.0
    total_tokens = 0
    for start in range(0, len(texts), batch_size):
        batch = texts[start : start + batch_size]
        enc = scorer_tok(
            batch,
            return_tensors="pt",
            return_token_type_ids=False,
            return_attention_mask=True,
            padding=True,
            truncation=True,
            max_length=max_length,
        ).to(device)
        input_ids = enc["input_ids"]
        if input_ids.size(1) < 2:
            continue
        logits = scorer(input_ids=input_ids, attention_mask=enc["attention_mask"]).logits.transpose(-1, -2)
        token_nll = F.cross_entropy(logits[..., :-1].float(), input_ids[..., 1:], reduction="none")
        if scorer_tok.eos_token_id is not None:
            first_eos = (input_ids == scorer_tok.eos_token_id).cumsum(-1) == 1
            token_mask = input_ids != scorer_tok.eos_token_id
            shift_mask = (first_eos[..., 1:] | token_mask[..., 1:]).contiguous()
        else:
            shift_mask = enc["attention_mask"][..., 1:].contiguous().bool()
        total_nll += float(token_nll[shift_mask].sum().detach().cpu())
        total_tokens += int(shift_mask.sum().detach().cpu())
    nll = total_nll / max(total_tokens, 1)
    return {"ppl": float(math.exp(min(nll, 50.0))), "nll_per_token": float(nll), "tokens": total_tokens}


def write_samples(path: Path, name: str, summary: dict, texts: list[str]) -> None:
    with path.open("w", encoding="utf-8") as f:
        f.write(name + "\n")
        f.write(json.dumps(summary, ensure_ascii=False, indent=2) + "\n\n")
        for i, text in enumerate(texts):
            f.write(f"===== sample {i:03d} =====\n{text}\n\n")


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--tokenizer_path", required=True)
    p.add_argument("--scorer", required=True)
    p.add_argument("--out_dir", required=True)
    p.add_argument("--n_samples", type=int, default=64)
    p.add_argument("--max_len", type=int, default=128)
    p.add_argument("--steps", type=int, default=512)
    p.add_argument("--decode_batch", type=int, default=16)
    p.add_argument("--score_batch", type=int, default=8)
    p.add_argument("--score_max_length", type=int, default=256)
    p.add_argument("--seed", type=int, default=20260507)
    args = p.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tokenizer = BpeTextTokenizer.from_file(args.tokenizer_path)
    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    model = build_model(ckpt, tokenizer, device)
    scorer_tok = AutoTokenizer.from_pretrained(args.scorer)
    if scorer_tok.pad_token_id is None:
        scorer_tok.pad_token = scorer_tok.eos_token
        scorer_tok.pad_token_id = scorer_tok.eos_token_id
    scorer = AutoModelForCausalLM.from_pretrained(args.scorer).to(device).eval()
    if getattr(scorer.config, "pad_token_id", None) is None:
        scorer.config.pad_token_id = scorer_tok.pad_token_id

    rows = []
    summary_jsonl = out_dir / "summary.jsonl"
    summary_jsonl.write_text("", encoding="utf-8")
    for cfg in default_configs():
        print(f"[decode] {cfg.name}", flush=True)
        ids, raw_texts, decode = decode_config(
            model,
            tokenizer,
            cfg,
            n_samples=args.n_samples,
            batch_size=args.decode_batch,
            max_len=args.max_len,
            steps=args.steps,
            seed=args.seed,
            device=device,
        )
        detok_texts = [lm1b_detokenizer(text) for text in raw_texts]
        detok_ppl = score_texts(
            [text for text in detok_texts if text.strip()],
            scorer,
            scorer_tok,
            batch_size=args.score_batch,
            max_length=args.score_max_length,
            device=device,
        )
        diversity = summarize_token_diversity(ids).__dict__
        summary = {
            "type": "summary",
            "name": cfg.name,
            "checkpoint": args.checkpoint,
            "step": ckpt.get("step"),
            "decode": decode,
            "detok_genppl": detok_ppl,
            "diversity": diversity,
        }
        write_samples(out_dir / f"{cfg.name}_raw_samples.txt", cfg.name, summary, raw_texts)
        write_samples(out_dir / f"{cfg.name}_detok_raw_samples.txt", cfg.name + "_detok", summary, detok_texts)
        with (out_dir / f"{cfg.name}_scored.jsonl").open("w", encoding="utf-8") as f:
            f.write(json.dumps(summary, ensure_ascii=False) + "\n")
            for i, (raw, detok) in enumerate(zip(raw_texts, detok_texts)):
                f.write(json.dumps({"type": "sample", "index": i, "raw_text": raw, "detok_text": detok}, ensure_ascii=False) + "\n")
        with summary_jsonl.open("a", encoding="utf-8") as f:
            f.write(json.dumps(summary, ensure_ascii=False) + "\n")
        row = {
            "name": cfg.name,
            "step": ckpt.get("step"),
            "detok_genppl": detok_ppl["ppl"],
            "sample_entropy": diversity["sample_entropy"],
            "distinct_2": diversity["distinct_2"],
            "top_token_mass": diversity["top_token_mass"],
            "model_t_mode": cfg.model_t_mode,
            "support_power": cfg.support_power,
            "semantic_power": cfg.semantic_power,
            "final_from": cfg.final_from,
            "endpoint_temp": cfg.endpoint_temp,
            "concentration_max": cfg.concentration_max,
            "update": cfg.update,
        }
        rows.append(row)
        print("[summary]", json.dumps(row, ensure_ascii=False), flush=True)

    keys = list(rows[0].keys())
    with (out_dir / "summary.tsv").open("w", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=keys, delimiter="\t")
        writer.writeheader()
        writer.writerows(rows)


if __name__ == "__main__":
    main()