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