| 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() |
|
|
