#!/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()