from __future__ import annotations import argparse import json import math import sys from dataclasses import asdict, is_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 model_time_for_step, sample_noise_simplex, state_for_model from flowtext_lab.genppl import filter_generated_texts, summarize_token_diversity from flowtext_lab.model import EndpointPredictor from flowtext_lab.tokenization import BpeTextTokenizer def extend_pos_embed(sd: dict[str, torch.Tensor], max_len: int, mode: str) -> dict[str, torch.Tensor]: sd = dict(sd) key = "pos_embed" if key not in sd: return sd pos = sd[key] old_len = int(pos.size(1)) if old_len == max_len: return sd if old_len > max_len: sd[key] = pos[:, :max_len].contiguous() return sd if mode == "repeat": reps = math.ceil(max_len / old_len) sd[key] = pos.repeat(1, reps, 1)[:, :max_len].contiguous() elif mode == "interpolate": x = pos.transpose(1, 2) y = F.interpolate(x, size=max_len, mode="linear", align_corners=True) sd[key] = y.transpose(1, 2).contiguous() else: raise ValueError(f"unknown pos_extend={mode}") return sd def build_model(ckpt: dict, tokenizer: BpeTextTokenizer, max_len: int, device: torch.device, pos_extend: str) -> EndpointPredictor: a = ckpt.get("args", {}) ckpt_state = ckpt["model"] if "output_bias" in a: output_bias = bool(a["output_bias"]) else: output_bias = "output_layer.linear.bias" in ckpt_state or "out_proj.bias" in ckpt_state vocab_size = int(a.get("effective_vocab_size", 0) or a.get("vocab_size_override", 0) or tokenizer.vocab_size) model = EndpointPredictor( vocab_size=vocab_size, max_len=max_len, d_model=int(a.get("d_model", 768)), n_heads=int(a.get("n_heads", 12)), n_layers=int(a.get("n_layers", 12)), dim_ff=int(a.get("dim_ff", 3072)), dropout=0.0, model_type=str(a.get("model_type", "ddit")), cond_dim=int(a.get("cond_dim", 128)), input_format=str(a.get("state_format", a.get("input_format", "prob"))), output_bias=output_bias, norm_type=str(a.get("norm_type", "layernorm")), ddit_mlp_type=str(a.get("ddit_mlp_type", "gelu")), ).to(device) state = extend_pos_embed(ckpt_state, max_len=max_len, mode=pos_extend) model.load_state_dict(state, strict=True) model.eval() return model def total_concentration(support_t: float, c_min: float, c_max: float) -> float: return math.exp(math.log(max(c_min, 1e-8)) + support_t * (math.log(max(c_max, c_min)) - math.log(max(c_min, 1e-8)))) def dirichlet_path_mean(endpoint: torch.Tensor, support_t: float, eps: float) -> torch.Tensor: vocab = endpoint.size(-1) mean = (1.0 - support_t) / float(vocab) + support_t * endpoint mean = mean.clamp_min(eps) return mean / mean.sum(dim=-1, keepdim=True).clamp_min(eps) def dirichlet_resample(mean: torch.Tensor, support_t: float, c_min: float, c_max: float, eps: float) -> torch.Tensor: conc = total_concentration(support_t, c_min, c_max) alpha = (mean * conc).clamp_min(eps) sample = torch._standard_gamma(alpha).clamp_min(eps) return sample / sample.sum(dim=-1, keepdim=True).clamp_min(eps) def current_anchor(probs: torch.Tensor, mode: str, eps: float) -> torch.Tensor: if mode == "onehot": return F.one_hot(probs.argmax(dim=-1), probs.size(-1)).to(dtype=probs.dtype) if mode == "sqrt_state": anchor = probs.clamp_min(eps).sqrt() else: anchor = probs.clamp_min(eps) return anchor / anchor.sum(dim=-1, keepdim=True).clamp_min(eps) def log_geodesic_mix(p: torch.Tensor, q: torch.Tensor, gamma: float, eps: float) -> torch.Tensor: log_mix = (1.0 - gamma) * p.clamp_min(eps).log() + gamma * q.clamp_min(eps).log() return torch.softmax(log_mix, dim=-1) def sqrt_geodesic_mix(p: torch.Tensor, q: torch.Tensor, gamma: float, eps: float) -> torch.Tensor: root = (1.0 - gamma) * p.clamp_min(eps).sqrt() + gamma * q.clamp_min(eps).sqrt() out = root.square().clamp_min(eps) return out / out.sum(dim=-1, keepdim=True).clamp_min(eps) def fisher_rao_mix(p: torch.Tensor, q: torch.Tensor, gamma: float, eps: float) -> torch.Tensor: a = p.clamp_min(eps).sqrt() b = q.clamp_min(eps).sqrt() dot = (a * b).sum(dim=-1, keepdim=True).clamp(-1.0 + 1e-6, 1.0 - 1e-6) theta = torch.acos(dot) sin_theta = torch.sin(theta).clamp_min(1e-6) left = torch.sin((1.0 - gamma) * theta) / sin_theta right = torch.sin(gamma * theta) / sin_theta root = left * a + right * b out = root.square().clamp_min(eps) return out / out.sum(dim=-1, keepdim=True).clamp_min(eps) def simplex_mix(p: torch.Tensor, q: torch.Tensor, gamma: float, eps: float, geometry: str) -> torch.Tensor: if geometry == "log": return log_geodesic_mix(p, q, gamma, eps) if geometry == "sqrt": return sqrt_geodesic_mix(p, q, gamma, eps) if geometry == "fisher": return fisher_rao_mix(p, q, gamma, eps) if geometry == "linear": out = (1.0 - gamma) * p + gamma * q out = out.clamp_min(eps) return out / out.sum(dim=-1, keepdim=True).clamp_min(eps) raise ValueError(geometry) def temperature(step: int, steps: int, early: float, late: float, temp_end: float, power: float) -> float: progress = step / max(steps, 1) if progress >= temp_end: return late rel = 1.0 - progress / max(temp_end, 1e-8) return late + (early - late) * (rel ** power) def make_time_grid( steps: int, *, schedule: str, logit_mean: float, logit_std: float, power: float, seed: int, device: torch.device, ) -> torch.Tensor: if steps <= 0: raise ValueError(f"steps must be positive, got {steps}") if schedule == "uniform": return torch.linspace(0.0, 1.0, steps + 1, device=device, dtype=torch.float32) if schedule == "logit_normal": if steps == 1: return torch.tensor([0.0, 1.0], device=device, dtype=torch.float32) generator = torch.Generator(device="cpu") generator.manual_seed(int(seed)) z = torch.randn((steps - 1,), generator=generator, dtype=torch.float32) middle = torch.sigmoid(z * float(logit_std) + float(logit_mean)).sort().values.to(device) return torch.cat( [ torch.zeros((1,), device=device, dtype=torch.float32), middle, torch.ones((1,), device=device, dtype=torch.float32), ] ) if schedule in {"power_low", "power_high"}: if steps == 1: return torch.tensor([0.0, 1.0], device=device, dtype=torch.float32) generator = torch.Generator(device="cpu") generator.manual_seed(int(seed)) u = torch.rand((steps - 1,), generator=generator, dtype=torch.float32) exponent = max(float(power), 1e-8) if schedule == "power_low": middle = u.pow(exponent) else: middle = 1.0 - (1.0 - u).pow(exponent) middle = middle.sort().values.to(device) return torch.cat( [ torch.zeros((1,), device=device, dtype=torch.float32), middle, torch.ones((1,), device=device, dtype=torch.float32), ] ) raise ValueError(f"unknown time schedule: {schedule}") def clamp_first_position(probs: torch.Tensor, first_ids: torch.Tensor | None) -> torch.Tensor: if first_ids is None: return probs probs = probs.clone() probs[:, 0, :].zero_() probs[:, 0, :].scatter_(1, first_ids[:, None], 1.0) return probs def final_decode_ids( probs: torch.Tensor, *, mode: str, temp: float, top_k: int, top_p: float, eps: float, ) -> torch.Tensor: if mode == "argmax": return probs.argmax(dim=-1) if mode != "sample": raise ValueError(mode) logits = probs.clamp_min(eps).log() / max(float(temp), eps) if top_k > 0 and top_k < logits.size(-1): kth = logits.topk(top_k, dim=-1).values[..., -1, None] logits = logits.masked_fill(logits < kth, -torch.inf) if 0.0 < top_p < 1.0: sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True) sorted_probs = F.softmax(sorted_logits, dim=-1) remove = sorted_probs.cumsum(dim=-1) > float(top_p) remove[..., 0] = False sorted_logits = sorted_logits.masked_fill(remove, -torch.inf) filtered = torch.full_like(logits, -torch.inf) logits = filtered.scatter(-1, sorted_idx, sorted_logits) sample_probs = F.softmax(logits, dim=-1) flat = sample_probs.reshape(-1, sample_probs.size(-1)) return torch.multinomial(flat, num_samples=1).view(probs.shape[:-1]) def soften_endpoint_with_prior( endpoint: torch.Tensor, t: float, *, mode: str, power: float, min_conf: float, max_conf: float, eps: float, ) -> tuple[torch.Tensor, float]: if mode == "none": return endpoint, 1.0 if mode != "uniform": raise ValueError(mode) alpha = float(min_conf) + (float(max_conf) - float(min_conf)) * (float(t) ** float(power)) alpha = max(0.0, min(1.0, alpha)) prior = 1.0 / float(endpoint.shape[-1]) softened = alpha * endpoint + (1.0 - alpha) * prior softened = softened.clamp_min(eps) softened = softened / softened.sum(dim=-1, keepdim=True).clamp_min(eps) return softened, alpha @torch.inference_mode() def decode( model: EndpointPredictor, tokenizer: BpeTextTokenizer, *, max_len: int, n_samples: int, batch_size: int, steps: int, seed: int, device: torch.device, decode_rule: str, support_power: float, semantic_power: float, early_temp: float, late_temp: float, temp_end: float, temp_power: float, hybrid_switch: float, tail_temp: float, c_min: float, c_max: float, model_t_mode: str, time_schedule: str, time_logit_mean: float, time_logit_std: float, time_power: float, input_noise_scale: float, input_noise_until: float, input_noise_dirichlet_concentration: float, endpoint_softening: str, endpoint_soft_power: float, endpoint_soft_min_conf: float, endpoint_soft_max_conf: float, final_from: str, final_decode: str, final_sample_temp: float, final_top_k: int, final_top_p: float, eps: float, fixed_first_token_id: int | None, fixed_first_initial_argmax: bool, ) -> tuple[list[list[int]], list[str], list[dict[str, object]]]: torch.manual_seed(seed) time_grid = make_time_grid( steps, schedule=time_schedule, logit_mean=time_logit_mean, logit_std=time_logit_std, power=time_power, seed=seed, device=device, ) all_ids: list[list[int]] = [] all_texts: list[str] = [] traces: list[dict[str, object]] = [] remaining = n_samples vocab_size = int(getattr(model, "vocab_size", tokenizer.vocab_size)) while remaining > 0: bs = min(batch_size, remaining) probs = sample_noise_simplex( (bs, max_len), vocab_size, device, eps, noise_mode="dirichlet", target_prob=1.0, noise_sigma=-1.0, dirichlet_concentration=1.0, ) fixed_first_ids: torch.Tensor | None = None if fixed_first_initial_argmax: fixed_first_ids = probs[:, 0, :].argmax(dim=-1) elif fixed_first_token_id is not None: fixed_first_ids = torch.full((bs,), int(fixed_first_token_id), dtype=torch.long, device=device) probs = clamp_first_position(probs, fixed_first_ids) attn = torch.ones((bs, max_len), dtype=torch.bool, device=device) last_endpoint = probs for step in range(steps): progress = float(time_grid[step].item()) next_progress = float(time_grid[step + 1].item()) dt = max(next_progress - progress, 0.0) if model_t_mode in {"pre", "flow"}: t = torch.full((bs,), float(progress), dtype=torch.float32, device=device) elif model_t_mode == "post": t = torch.full((bs,), float(next_progress), dtype=torch.float32, device=device) else: t = model_time_for_step(model_t_mode, step, steps, bs, device, dtype=torch.float32) temp = temperature(step, steps, early_temp, late_temp, temp_end, temp_power) if tail_temp > 0 and progress >= hybrid_switch: temp = tail_temp model_probs = probs if input_noise_scale > 0.0 and progress < input_noise_until: fresh_noise = sample_noise_simplex( (bs, max_len), vocab_size, device, eps, noise_mode="dirichlet", target_prob=1.0, noise_sigma=-1.0, dirichlet_concentration=input_noise_dirichlet_concentration, ) noisy = progress * probs + (1.0 - progress) * float(input_noise_scale) * fresh_noise model_probs = noisy.clamp_min(eps) model_probs = model_probs / model_probs.sum(dim=-1, keepdim=True).clamp_min(eps) logits = model(state_for_model(model, model_probs, eps), t, attn).float() raw_endpoint = F.softmax(logits / temp, dim=-1) endpoint, endpoint_alpha = soften_endpoint_with_prior( raw_endpoint, next_progress, mode=endpoint_softening, power=endpoint_soft_power, min_conf=endpoint_soft_min_conf, max_conf=endpoint_soft_max_conf, eps=eps, ) last_endpoint = endpoint support_t = next_progress ** support_power if decode_rule == "dirichlet_resample": probs = dirichlet_resample(dirichlet_path_mean(endpoint, support_t, eps), support_t, c_min, c_max, eps) elif decode_rule == "dual_line_resample": semantic_t = next_progress ** semantic_power anchor = current_anchor(probs, "state", 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) probs = dirichlet_resample(dirichlet_path_mean(forward_endpoint, support_t, eps), support_t, c_min, c_max, eps) elif decode_rule == "dual_replace_resample": semantic_t = next_progress ** semantic_power anchor = current_anchor(probs, "state", eps) replace = torch.rand((bs, max_len, 1), device=device) < semantic_t forward_endpoint = torch.where(replace, endpoint, anchor) forward_endpoint = forward_endpoint.clamp_min(eps) forward_endpoint = forward_endpoint / forward_endpoint.sum(dim=-1, keepdim=True).clamp_min(eps) probs = dirichlet_resample(dirichlet_path_mean(forward_endpoint, support_t, eps), support_t, c_min, c_max, eps) elif decode_rule in {"log_dual_resample", "sqrt_dual_resample", "fisher_dual_resample"}: geometry = decode_rule.split("_", 1)[0] semantic_t = next_progress ** semantic_power anchor = current_anchor(probs, "state", eps) forward_endpoint = simplex_mix(anchor, endpoint, semantic_t, eps, geometry) probs = dirichlet_resample(dirichlet_path_mean(forward_endpoint, support_t, eps), support_t, c_min, c_max, eps) elif decode_rule == "flowmap": gamma = min(dt / max(1.0 - progress, eps), 1.0) probs = probs + gamma * (endpoint - probs) probs = probs.clamp_min(eps) probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(eps) elif decode_rule in {"log_geodesic", "sqrt_geodesic", "fisher_geodesic"}: geometry = decode_rule.split("_", 1)[0] gamma = min(dt / max(1.0 - progress, eps), 1.0) probs = simplex_mix(probs, endpoint, gamma, eps, geometry) elif decode_rule in {"hybrid_log_flowmap", "hybrid_log_dirres", "hybrid_log_logflow"}: if progress < hybrid_switch: local = min(1.0, next_progress / max(hybrid_switch, 1e-8)) semantic_t = local ** semantic_power anchor = current_anchor(probs, "state", eps) forward_endpoint = simplex_mix(anchor, endpoint, semantic_t, eps, "log") probs = dirichlet_resample(dirichlet_path_mean(forward_endpoint, support_t, eps), support_t, c_min, c_max, eps) elif decode_rule == "hybrid_log_flowmap": gamma = min(dt / max(1.0 - progress, eps), 1.0) probs = simplex_mix(probs, endpoint, gamma, eps, "linear") elif decode_rule == "hybrid_log_logflow": gamma = min(dt / max(1.0 - progress, eps), 1.0) probs = simplex_mix(probs, endpoint, gamma, eps, "log") else: probs = dirichlet_resample(dirichlet_path_mean(endpoint, support_t, eps), support_t, c_min, c_max, eps) else: raise ValueError(decode_rule) probs = clamp_first_position(probs, fixed_first_ids) if step in {0, 1, 3, 7, 15, 31, 63, steps - 1}: ids0 = probs.argmax(dim=-1)[0].detach().cpu().tolist() raw_maxprob = raw_endpoint[0].amax(dim=-1).mean().detach().item() soft_maxprob = endpoint[0].amax(dim=-1).mean().detach().item() traces.append({ "step": step + 1, "progress": progress, "next_progress": next_progress, "dt": dt, "temperature": temp, "endpoint_alpha": endpoint_alpha, "raw_endpoint_mean_maxprob": raw_maxprob, "effective_endpoint_mean_maxprob": soft_maxprob, "sample0_text": tokenizer.decode(ids0, stop_at_eos=False, skip_special_tokens=False)[:480], }) if final_from == "state": final = probs elif final_from == "endpoint": final = last_endpoint elif final_from == "blend": final = 0.5 * probs + 0.5 * last_endpoint else: raise ValueError(final_from) final = clamp_first_position(final, fixed_first_ids) ids_tensor = final_decode_ids( final, mode=final_decode, temp=final_sample_temp, top_k=final_top_k, top_p=final_top_p, eps=eps, ) ids = ids_tensor.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 print(f"[decode] max_len={max_len} generated={n_samples-remaining}/{n_samples}", flush=True) return all_ids, all_texts, traces def score_with_gpt2(texts: list[str], scorer_path: str, batch_size: int, max_length: int, device: torch.device) -> dict[str, float]: scorer_tok = AutoTokenizer.from_pretrained(scorer_path, local_files_only=True) if scorer_tok.pad_token is None: scorer_tok.pad_token = scorer_tok.eos_token scorer = AutoModelForCausalLM.from_pretrained(scorer_path, local_files_only=True).to(device).eval() 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", padding=True, truncation=True, max_length=max_length).to(device) input_ids = enc["input_ids"] attn = enc["attention_mask"] if input_ids.size(1) < 2: continue logits = scorer(input_ids=input_ids, attention_mask=attn).logits.transpose(-1, -2) nll = F.cross_entropy(logits[..., :-1].float(), input_ids[..., 1:], reduction="none") mask = attn[..., 1:].bool() total_nll += float(nll[mask].sum().detach().cpu()) total_tokens += int(mask.sum().detach().cpu()) del scorer if device.type == "cuda": torch.cuda.empty_cache() mean = total_nll / max(total_tokens, 1) return {"gen_ppl": math.exp(min(20.0, mean)), "gen_nll": mean, "gen_tokens": total_tokens} def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--checkpoint", required=True) ap.add_argument("--tokenizer_path", required=True) ap.add_argument("--out_dir", required=True) ap.add_argument("--max_lens", default="128,1024") ap.add_argument("--n_samples", type=int, default=16) ap.add_argument("--batch_size", type=int, default=2) ap.add_argument("--steps", type=int, default=128) ap.add_argument( "--decode_rule", choices=[ "dual_line_resample", "dual_replace_resample", "dirichlet_resample", "flowmap", "log_dual_resample", "sqrt_dual_resample", "fisher_dual_resample", "log_geodesic", "sqrt_geodesic", "fisher_geodesic", "hybrid_log_flowmap", "hybrid_log_dirres", "hybrid_log_logflow", ], default="dual_line_resample", ) ap.add_argument("--pos_extend", choices=["repeat", "interpolate"], default="repeat") ap.add_argument("--support_power", type=float, default=1.0) ap.add_argument("--semantic_power", type=float, default=1.5) ap.add_argument("--early_temp", type=float, default=2.8) ap.add_argument("--late_temp", type=float, default=1.45) ap.add_argument("--temp_end", type=float, default=0.55) ap.add_argument("--temp_power", type=float, default=1.5) ap.add_argument("--hybrid_switch", type=float, default=0.5) ap.add_argument("--tail_temp", type=float, default=-1.0) ap.add_argument("--c_min", type=float, default=1.0) ap.add_argument("--c_max", type=float, default=1024.0) ap.add_argument( "--model_t_mode", choices=["pre", "post", "flow", "linear", "const0", "const05", "const1", "random"], default="flow", ) ap.add_argument("--time_schedule", choices=["uniform", "logit_normal", "power_low", "power_high"], default="uniform") ap.add_argument("--time_logit_mean", type=float, default=-1.5) ap.add_argument("--time_logit_std", type=float, default=0.8) ap.add_argument("--time_power", type=float, default=2.0) ap.add_argument("--input_noise_scale", type=float, default=0.0) ap.add_argument("--input_noise_until", type=float, default=1.0) ap.add_argument("--input_noise_dirichlet_concentration", type=float, default=1.0) ap.add_argument("--endpoint_softening", choices=["none", "uniform"], default="none") ap.add_argument("--endpoint_soft_power", type=float, default=2.0) ap.add_argument("--endpoint_soft_min_conf", type=float, default=0.0) ap.add_argument("--endpoint_soft_max_conf", type=float, default=1.0) ap.add_argument("--final_from", choices=["state", "endpoint", "blend"], default="blend") ap.add_argument("--final_decode", choices=["argmax", "sample"], default="argmax") ap.add_argument("--final_sample_temp", type=float, default=1.0) ap.add_argument("--final_top_k", type=int, default=0) ap.add_argument("--final_top_p", type=float, default=1.0) ap.add_argument("--fixed_first_token_id", type=int, default=-1) ap.add_argument("--fixed_first_token_text", default="") ap.add_argument("--fixed_first_initial_argmax", action="store_true") ap.add_argument("--scorer", default="/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard") ap.add_argument("--score", action="store_true") ap.add_argument("--use_ema", action="store_true", help="Use ema_model from checkpoint if present.") ap.add_argument("--seed", type=int, default=20260514) args = ap.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") tok = BpeTextTokenizer.from_file(args.tokenizer_path) ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False, mmap=True) if args.use_ema and "ema_model" in ckpt: ckpt = dict(ckpt) ckpt["model"] = ckpt["ema_model"] out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) fixed_first_token_id: int | None = None if args.fixed_first_token_text: encoded = tok.encode(args.fixed_first_token_text, add_eos=False, add_special_tokens=False) if not encoded: raise ValueError(f"fixed_first_token_text encoded to no tokens: {args.fixed_first_token_text!r}") fixed_first_token_id = int(encoded[0]) elif args.fixed_first_token_id >= 0: fixed_first_token_id = int(args.fixed_first_token_id) summary = [] for max_len_s in args.max_lens.split(","): max_len = int(max_len_s) model = build_model(ckpt, tok, max_len, device, args.pos_extend) ids, texts, traces = decode( model, tok, max_len=max_len, n_samples=args.n_samples, batch_size=args.batch_size, steps=args.steps, seed=args.seed + max_len, device=device, decode_rule=args.decode_rule, support_power=args.support_power, semantic_power=args.semantic_power, early_temp=args.early_temp, late_temp=args.late_temp, temp_end=args.temp_end, temp_power=args.temp_power, hybrid_switch=args.hybrid_switch, tail_temp=args.tail_temp, c_min=args.c_min, c_max=args.c_max, model_t_mode=args.model_t_mode, time_schedule=args.time_schedule, time_logit_mean=args.time_logit_mean, time_logit_std=args.time_logit_std, time_power=args.time_power, input_noise_scale=args.input_noise_scale, input_noise_until=args.input_noise_until, input_noise_dirichlet_concentration=args.input_noise_dirichlet_concentration, endpoint_softening=args.endpoint_softening, endpoint_soft_power=args.endpoint_soft_power, endpoint_soft_min_conf=args.endpoint_soft_min_conf, endpoint_soft_max_conf=args.endpoint_soft_max_conf, final_from=args.final_from, final_decode=args.final_decode, final_sample_temp=args.final_sample_temp, final_top_k=args.final_top_k, final_top_p=args.final_top_p, eps=1e-8, fixed_first_token_id=fixed_first_token_id, fixed_first_initial_argmax=args.fixed_first_initial_argmax, ) filt_result = filter_generated_texts(texts, min_chars=0, normalize_whitespace=True, drop_empty=False) filt = filt_result[0] if isinstance(filt_result, tuple) else filt_result diversity_result = summarize_token_diversity(ids) diversity = asdict(diversity_result) if is_dataclass(diversity_result) else dict(diversity_result) rec = { "checkpoint": args.checkpoint, "ckpt_step": int(ckpt.get("step", -1)), "max_len": max_len, "decode_rule": args.decode_rule, "support_power": args.support_power, "semantic_power": args.semantic_power, "steps": args.steps, "c_min": args.c_min, "c_max": args.c_max, "model_t_mode": args.model_t_mode, "time_schedule": args.time_schedule, "time_logit_mean": args.time_logit_mean, "time_logit_std": args.time_logit_std, "time_power": args.time_power, "input_noise_scale": args.input_noise_scale, "input_noise_until": args.input_noise_until, "input_noise_dirichlet_concentration": args.input_noise_dirichlet_concentration, "endpoint_softening": args.endpoint_softening, "endpoint_soft_power": args.endpoint_soft_power, "endpoint_soft_min_conf": args.endpoint_soft_min_conf, "endpoint_soft_max_conf": args.endpoint_soft_max_conf, "final_from": args.final_from, "final_decode": args.final_decode, "final_sample_temp": args.final_sample_temp, "final_top_k": args.final_top_k, "final_top_p": args.final_top_p, "early_temp": args.early_temp, "late_temp": args.late_temp, "temp_end": args.temp_end, "temp_power": args.temp_power, "pos_extend": args.pos_extend, "fixed_first_token_id": fixed_first_token_id, "fixed_first_token_text": args.fixed_first_token_text, "fixed_first_initial_argmax": bool(args.fixed_first_initial_argmax), "use_ema": bool(args.use_ema and "ema_model" in ckpt), "n_samples": len(texts), **diversity, "texts_preview": filt[:4], } if args.score: rec.update(score_with_gpt2(filt, args.scorer, batch_size=2, max_length=min(max_len, 1024), device=device)) (out_dir / f"context{max_len}_samples.txt").write_text("\n\n---\n\n".join(filt), encoding="utf-8") (out_dir / f"context{max_len}_trace.json").write_text(json.dumps(traces, ensure_ascii=False, indent=2), encoding="utf-8") summary.append(rec) del model if device.type == "cuda": torch.cuda.empty_cache() (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True) if __name__ == "__main__": main()