#!/usr/bin/env python3 from __future__ import annotations import argparse import collections import json import math import re import sys import time from pathlib import Path from typing import Any import torch import torch.nn.functional as F 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.bridges import smooth_onehot from flowtext_lab.decode import ( model_time_for_step, sample_noise_simplex, state_for_model, ) from flowtext_lab.model import EndpointPredictor from flowtext_lab.tokenization import BpeTextTokenizer from train import ELFConditionalCollator, load_elf_conditional_dataset def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser( description=( "Conditional WMT14 de-en generation eval for the PyTorch LTA/ELF model. " "Writes generated translations and official-style BLEU/ROUGE metrics." ) ) p.add_argument("--checkpoint", required=True) p.add_argument("--data_path", default="", help="Validation dataset path. Defaults to ckpt args eval_data_path.") p.add_argument("--tokenizer_path", default="", help="Defaults to ckpt args tokenizer_path.") p.add_argument("--output_dir", default="", help="Defaults to /wmt14_cond_eval.") p.add_argument("--num_samples", type=int, default=128) p.add_argument("--batch_size", type=int, default=16) p.add_argument("--max_len", type=int, default=0, help="Defaults to ckpt args max_len.") p.add_argument("--max_input_len", type=int, default=0, help="Defaults to ckpt args max_input_len.") p.add_argument("--conditional_pad_token", choices=["auto", "pad", "eos"], default="auto") p.add_argument("--use_ema", action="store_true", help="Use ema_model from checkpoint if present.") p.add_argument("--infer_steps", type=int, default=0, help="Defaults to ckpt args infer_steps.") p.add_argument("--decode_damping", type=float, default=math.nan, help="Defaults to ckpt args decode_damping.") p.add_argument("--max_gamma", type=float, default=math.nan, help="Defaults to ckpt args max_gamma or 1.0.") p.add_argument("--decode_solver", choices=["auto", "simple", "flowmap"], default="auto") p.add_argument( "--decode_rule", choices=[ "auto", "simple", "flowmap", "dirichlet_resample", "dual_line_resample", "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="auto", help="LM1B/OWT-style decode rule. auto falls back to --decode_solver/ckpt decode_solver.", ) p.add_argument( "--model_t_mode", choices=["pre", "post", "linear", "flow", "const0", "const05", "const1", "random"], default="flow", ) p.add_argument("--temperature", type=float, default=1.0) p.add_argument("--early_temp", type=float, default=2.8) p.add_argument("--late_temp", type=float, default=1.45) p.add_argument("--temp_end", type=float, default=0.55) p.add_argument("--temp_power", type=float, default=1.5) p.add_argument("--tail_temp", type=float, default=-1.0) p.add_argument("--support_power", type=float, default=1.0) p.add_argument("--semantic_power", type=float, default=1.5) p.add_argument("--hybrid_switch", type=float, default=0.5) p.add_argument("--c_min", type=float, default=1.0) p.add_argument("--c_max", type=float, default=1024.0) p.add_argument( "--final_endpoint_blend", type=float, default=-1.0, help="Blend final state with last endpoint before argmax. <0 uses 0.5 for LM1B resample/geodesic rules, 0 for simple/flowmap.", ) p.add_argument("--noise_init", choices=["auto", "uniform", "logistic_normal", "dirichlet"], default="auto") p.add_argument("--noise_sigma", type=float, default=math.nan) p.add_argument("--dirichlet_init_concentration", type=float, default=math.nan) p.add_argument("--target_prob", type=float, default=math.nan) p.add_argument("--eps", type=float, default=1e-8) p.add_argument("--seed", type=int, default=1234) p.add_argument("--detokenizer", default=None) p.add_argument("--print_samples", type=int, default=3) return p.parse_args() def arg_or_ckpt(args: argparse.Namespace, train_args: dict[str, Any], name: str, default: Any) -> Any: value = getattr(args, name) if isinstance(value, float) and math.isnan(value): return train_args.get(name, default) if isinstance(value, int) and value == 0: return train_args.get(name, default) if isinstance(value, str) and value == "auto": return train_args.get(name, default) if isinstance(value, str) and value == "": return train_args.get(name, default) return value def build_model_from_ckpt( ckpt: dict[str, Any], tokenizer: BpeTextTokenizer, max_len: int, device: torch.device, *, use_ema: bool, ) -> EndpointPredictor: train_args = ckpt.get("args", {}) state_key = "ema_model" if use_ema and "ema_model" in ckpt else "model" state = ckpt[state_key] output_bias = bool(train_args.get("output_bias", "output_layer.linear.bias" in state or "out_proj.bias" in state)) model = EndpointPredictor( vocab_size=int(ckpt.get("vocab_size", tokenizer.vocab_size) or tokenizer.vocab_size), max_len=max_len, d_model=int(train_args.get("d_model", 384)), n_heads=int(train_args.get("n_heads", 6)), n_layers=int(train_args.get("n_layers", 6)), dim_ff=int(train_args.get("dim_ff", 1536)), dropout=0.0, model_type=str(train_args.get("model_type", "transformer")), cond_dim=int(train_args.get("cond_dim", 128)), input_format=str(train_args.get("state_format", train_args.get("input_format", "logprob"))), output_bias=output_bias, norm_type=str(train_args.get("norm_type", "layernorm")), elf_num_time_tokens=int(train_args.get("elf_num_time_tokens", 4)), elf_num_model_mode_tokens=int(train_args.get("elf_num_model_mode_tokens", 0)), qk_norm=bool(train_args.get("qk_norm", True)), output_init_std=float(train_args.get("output_init_std", -1.0)), ).to(device) 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(float(c_min), 1e-8)) + float(support_t) * (math.log(max(float(c_max), float(c_min))) - math.log(max(float(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 - float(support_t)) / float(vocab) + float(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) root = (torch.sin((1.0 - gamma) * theta) / sin_theta) * a + (torch.sin(gamma * theta) / sin_theta) * 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 scheduled_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 float(late) rel = 1.0 - progress / max(temp_end, 1e-8) return float(late) + (float(early) - float(late)) * (rel ** float(power)) def apply_lm1b_decode_update( *, decode_rule: str, probs: torch.Tensor, endpoint: torch.Tensor, step: int, steps: int, support_power: float, semantic_power: float, hybrid_switch: float, c_min: float, c_max: float, eps: float, ) -> torch.Tensor: progress = step / max(steps, 1) next_progress = (step + 1) / max(steps, 1) support_t = next_progress ** float(support_power) if decode_rule == "dirichlet_resample": return dirichlet_resample(dirichlet_path_mean(endpoint, support_t, eps), support_t, c_min, c_max, eps) if decode_rule == "dual_line_resample": semantic_t = next_progress ** float(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) return dirichlet_resample(dirichlet_path_mean(forward_endpoint, support_t, eps), support_t, c_min, c_max, eps) if decode_rule in {"log_dual_resample", "sqrt_dual_resample", "fisher_dual_resample"}: geometry = decode_rule.split("_", 1)[0] semantic_t = next_progress ** float(semantic_power) anchor = current_anchor(probs, "state", eps) forward_endpoint = simplex_mix(anchor, endpoint, semantic_t, eps, geometry) return dirichlet_resample(dirichlet_path_mean(forward_endpoint, support_t, eps), support_t, c_min, c_max, eps) if decode_rule == "flowmap": gamma = min(1.0 / max(steps - step, 1), 1.0) out = probs + gamma * (endpoint - probs) out = out.clamp_min(eps) return out / out.sum(dim=-1, keepdim=True).clamp_min(eps) if decode_rule in {"log_geodesic", "sqrt_geodesic", "fisher_geodesic"}: geometry = decode_rule.split("_", 1)[0] gamma = min(1.0 / max(steps - step, 1), 1.0) return simplex_mix(probs, endpoint, gamma, eps, geometry) if decode_rule in {"hybrid_log_flowmap", "hybrid_log_dirres", "hybrid_log_logflow"}: if progress < hybrid_switch: local = min(1.0, next_progress / max(float(hybrid_switch), 1e-8)) semantic_t = local ** float(semantic_power) anchor = current_anchor(probs, "state", eps) forward_endpoint = simplex_mix(anchor, endpoint, semantic_t, eps, "log") return dirichlet_resample(dirichlet_path_mean(forward_endpoint, support_t, eps), support_t, c_min, c_max, eps) if decode_rule == "hybrid_log_flowmap": gamma = min(1.0 / max(steps - step, 1), 1.0) return simplex_mix(probs, endpoint, gamma, eps, "linear") if decode_rule == "hybrid_log_logflow": gamma = min(1.0 / max(steps - step, 1), 1.0) return simplex_mix(probs, endpoint, gamma, eps, "log") return dirichlet_resample(dirichlet_path_mean(endpoint, support_t, eps), support_t, c_min, c_max, eps) raise ValueError(decode_rule) @torch.no_grad() def conditional_decode( model: EndpointPredictor, init_probs: torch.Tensor, attn_mask: torch.Tensor, prefix_lock: torch.Tensor, prefix_probs: torch.Tensor, *, infer_steps: int, damping: float, max_gamma: float, eps: float, solver: str, decode_rule: str, model_t_mode: str, temperature: float, early_temp: float, late_temp: float, temp_end: float, temp_power: float, tail_temp: float, support_power: float, semantic_power: float, hybrid_switch: float, c_min: float, c_max: float, final_endpoint_blend: float, pad_id: int, ) -> torch.Tensor: probs = init_probs.float().clone() temp = max(float(temperature), eps) pad_fill = torch.zeros_like(probs) pad_fill[..., int(pad_id)] = 1.0 logits = None last_endpoint = probs for step in range(int(infer_steps)): progress = step / max(infer_steps, 1) if decode_rule in {"auto", "simple", "flowmap"} and solver == "flowmap": s = step / max(infer_steps, 1) t_next = (step + 1) / max(infer_steps, 1) gamma = float(damping) * ((t_next - s) / max(1.0 - s, eps)) if max_gamma > 0: gamma = min(gamma, float(max_gamma)) elif decode_rule in {"auto", "simple", "flowmap"} and solver == "simple": gamma = min(float(max_gamma), float(damping) / max(infer_steps, 1)) if model_t_mode == "pre": t = torch.full( (probs.size(0),), float(step / max(infer_steps, 1)), device=probs.device, dtype=torch.float32, ) elif model_t_mode == "post": t = torch.full( (probs.size(0),), float((step + 1) / max(infer_steps, 1)), device=probs.device, dtype=torch.float32, ) else: t = model_time_for_step( model_t_mode, step, infer_steps, probs.size(0), probs.device, dtype=torch.float32, ) logits = model(state_for_model(model, probs, eps), t, attn_mask) if decode_rule in {"auto", "simple", "flowmap"}: temp = max(float(temperature), eps) else: temp = scheduled_temperature(step, infer_steps, early_temp, late_temp, temp_end, temp_power) if tail_temp > 0 and progress >= hybrid_switch: temp = float(tail_temp) endpoint = F.softmax(logits.float() / temp, dim=-1) last_endpoint = endpoint if decode_rule in {"auto", "simple", "flowmap"}: probs = probs + gamma * (endpoint - probs) probs = probs.clamp_min(eps) probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(eps) else: probs = apply_lm1b_decode_update( decode_rule=decode_rule, probs=probs, endpoint=endpoint, step=step, steps=infer_steps, support_power=support_power, semantic_power=semantic_power, hybrid_switch=hybrid_switch, c_min=c_min, c_max=c_max, eps=eps, ) probs = torch.where(prefix_lock.unsqueeze(-1), prefix_probs, probs) probs = torch.where(attn_mask.unsqueeze(-1), probs, pad_fill) if logits is None: logits = model(state_for_model(model, probs, eps), torch.zeros(probs.size(0), device=probs.device), attn_mask) if final_endpoint_blend < 0: final_endpoint_blend = 0.0 if decode_rule in {"auto", "simple", "flowmap"} else 0.5 if final_endpoint_blend > 0: blend = min(max(float(final_endpoint_blend), 0.0), 1.0) probs = (1.0 - blend) * probs + blend * last_endpoint probs = probs.clamp_min(eps) probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(eps) probs = torch.where(prefix_lock.unsqueeze(-1), prefix_probs, probs) probs = torch.where(attn_mask.unsqueeze(-1), probs, pad_fill) return probs def compute_text_metrics(hypotheses: list[str], references: list[str]) -> dict[str, Any]: metrics: dict[str, Any] = { "num_samples": len(hypotheses), "empty_rate": sum(1 for x in hypotheses if not x.strip()) / max(len(hypotheses), 1), "exact_match": sum(h.strip() == r.strip() for h, r in zip(hypotheses, references)) / max(len(hypotheses), 1), "avg_generated_chars": sum(len(x) for x in hypotheses) / max(len(hypotheses), 1), "avg_reference_chars": sum(len(x) for x in references) / max(len(references), 1), } warnings: list[str] = [] try: import sacrebleu metrics["bleu"] = sacrebleu.corpus_bleu( hypotheses, [references], lowercase=True, use_effective_order=True, ).score metrics["chrf"] = sacrebleu.corpus_chrf(hypotheses, [references]).score metrics["bleu_backend"] = "sacrebleu" metrics["chrf_backend"] = "sacrebleu" except Exception as exc: # pragma: no cover - depends on remote env metrics["bleu"] = fallback_corpus_bleu(hypotheses, references) metrics["chrf"] = fallback_corpus_chrf(hypotheses, references) metrics["bleu_backend"] = "fallback_whitespace_bleu" metrics["chrf_backend"] = "fallback_char_f" warnings.append(f"sacrebleu unavailable, used fallback metrics: {exc}") try: from rouge_score import rouge_scorer scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True) r1: list[float] = [] r2: list[float] = [] rl: list[float] = [] for hyp, ref in zip(hypotheses, references): score = scorer.score(ref, hyp) r1.append(score["rouge1"].fmeasure * 100.0) r2.append(score["rouge2"].fmeasure * 100.0) rl.append(score["rougeL"].fmeasure * 100.0) metrics["rouge1"] = sum(r1) / max(len(r1), 1) metrics["rouge2"] = sum(r2) / max(len(r2), 1) metrics["rougeL"] = sum(rl) / max(len(rl), 1) metrics["rouge_backend"] = "rouge_score" except Exception as exc: # pragma: no cover - depends on remote env rouge = fallback_rouge(hypotheses, references) metrics.update(rouge) metrics["rouge_backend"] = "fallback_whitespace_lcs_no_stem" warnings.append(f"rouge_score unavailable, used fallback metrics: {exc}") if warnings: metrics["metric_warnings"] = warnings return metrics _TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE) def metric_tokens(text: str) -> list[str]: return _TOKEN_RE.findall(text.lower()) def ngram_counts(tokens: list[str], n: int) -> collections.Counter[tuple[str, ...]]: if len(tokens) < n: return collections.Counter() return collections.Counter(tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)) def fallback_corpus_bleu(hypotheses: list[str], references: list[str]) -> float: """Small corpus BLEU fallback for offline clusters. It is intentionally labeled as a fallback because official ELF uses sacrebleu's 13a tokenization. This keeps smoke/debug eval numerically populated when the cluster cannot install sacrebleu. """ hyp_tokens = [metric_tokens(x) for x in hypotheses] ref_tokens = [metric_tokens(x) for x in references] hyp_len = sum(len(x) for x in hyp_tokens) ref_len = sum(len(x) for x in ref_tokens) if hyp_len == 0: return 0.0 precisions = [] smooth = 1.0 for n in range(1, 5): clipped = 0 total = 0 for hyp, ref in zip(hyp_tokens, ref_tokens): hyp_ngrams = ngram_counts(hyp, n) ref_ngrams = ngram_counts(ref, n) total += sum(hyp_ngrams.values()) clipped += sum(min(count, ref_ngrams[ng]) for ng, count in hyp_ngrams.items()) if total == 0: continue if clipped == 0: smooth *= 2.0 precisions.append(1.0 / (smooth * total)) else: precisions.append(clipped / total) if not precisions: return 0.0 bp = 1.0 if hyp_len > ref_len else math.exp(1.0 - (ref_len / max(hyp_len, 1))) return 100.0 * bp * math.exp(sum(math.log(p) for p in precisions) / len(precisions)) def fallback_corpus_chrf(hypotheses: list[str], references: list[str], beta: float = 2.0) -> float: hyp = "\n".join(hypotheses) ref = "\n".join(references) scores = [] for n in range(1, 7): h = ngram_counts(list(hyp), n) r = ngram_counts(list(ref), n) overlap = sum(min(count, r[ng]) for ng, count in h.items()) prec = overlap / max(sum(h.values()), 1) rec = overlap / max(sum(r.values()), 1) denom = (beta * beta * prec) + rec scores.append(0.0 if denom == 0.0 else (1 + beta * beta) * prec * rec / denom) return 100.0 * sum(scores) / max(len(scores), 1) def lcs_len(a: list[str], b: list[str]) -> int: if not a or not b: return 0 prev = [0] * (len(b) + 1) for tok_a in a: cur = [0] for j, tok_b in enumerate(b, start=1): cur.append(prev[j - 1] + 1 if tok_a == tok_b else max(prev[j], cur[-1])) prev = cur return prev[-1] def f_measure(overlap: int, pred_total: int, ref_total: int) -> float: if overlap <= 0 or pred_total <= 0 or ref_total <= 0: return 0.0 precision = overlap / pred_total recall = overlap / ref_total return 100.0 * (2.0 * precision * recall) / max(precision + recall, 1e-12) def fallback_rouge(hypotheses: list[str], references: list[str]) -> dict[str, float]: r1: list[float] = [] r2: list[float] = [] rl: list[float] = [] for hyp_text, ref_text in zip(hypotheses, references): hyp = metric_tokens(hyp_text) ref = metric_tokens(ref_text) hyp1 = ngram_counts(hyp, 1) ref1 = ngram_counts(ref, 1) hyp2 = ngram_counts(hyp, 2) ref2 = ngram_counts(ref, 2) overlap1 = sum(min(count, ref1[ng]) for ng, count in hyp1.items()) overlap2 = sum(min(count, ref2[ng]) for ng, count in hyp2.items()) r1.append(f_measure(overlap1, sum(hyp1.values()), sum(ref1.values()))) r2.append(f_measure(overlap2, sum(hyp2.values()), sum(ref2.values()))) rl.append(f_measure(lcs_len(hyp, ref), len(hyp), len(ref))) return { "rouge1": sum(r1) / max(len(r1), 1), "rouge2": sum(r2) / max(len(r2), 1), "rougeL": sum(rl) / max(len(rl), 1), } def decode_target_segment( tokenizer: BpeTextTokenizer, ids: torch.Tensor, cond_len: int, gen_len: int, *, detokenizer: str | None, ) -> str: segment = ids[int(cond_len) : int(cond_len) + int(gen_len)].detach().cpu().tolist() return tokenizer.decode( segment, stop_at_eos=True, skip_special_tokens=True, detokenizer=detokenizer, ) @torch.no_grad() def main() -> None: args = parse_args() torch.manual_seed(int(args.seed)) if torch.cuda.is_available(): torch.cuda.manual_seed_all(int(args.seed)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") ckpt = torch.load(args.checkpoint, map_location="cpu") train_args = ckpt.get("args", {}) tokenizer_path = arg_or_ckpt(args, train_args, "tokenizer_path", "") data_path = arg_or_ckpt(args, train_args, "data_path", train_args.get("eval_data_path", "")) if args.data_path: data_path = args.data_path elif train_args.get("eval_data_path"): data_path = train_args["eval_data_path"] if not tokenizer_path: raise ValueError("--tokenizer_path is required when checkpoint args do not contain tokenizer_path") if not data_path: raise ValueError("--data_path is required when checkpoint args do not contain eval_data_path") tokenizer = BpeTextTokenizer.from_file(tokenizer_path) max_len = int(arg_or_ckpt(args, train_args, "max_len", 128)) max_input_len = int(arg_or_ckpt(args, train_args, "max_input_len", min(64, max_len))) infer_steps = int(arg_or_ckpt(args, train_args, "infer_steps", 128)) damping = float(arg_or_ckpt(args, train_args, "decode_damping", 1.0)) max_gamma = float(arg_or_ckpt(args, train_args, "max_gamma", 1.0)) solver = str(train_args.get("decode_solver", "flowmap") if args.decode_solver == "auto" else args.decode_solver) decode_rule = solver if args.decode_rule == "auto" else str(args.decode_rule) if decode_rule in {"simple", "flowmap"}: solver = decode_rule noise_mode = str(train_args.get("noise_init", "dirichlet") if args.noise_init == "auto" else args.noise_init) if noise_mode == "logistic_normal" and float(train_args.get("target_prob", 0.99)) >= 1.0 and math.isnan(args.target_prob): target_prob = 0.99 else: target_prob = float(arg_or_ckpt(args, train_args, "target_prob", 0.99)) noise_sigma = float(arg_or_ckpt(args, train_args, "noise_sigma", -1.0)) dirichlet_conc = float( arg_or_ckpt( args, train_args, "dirichlet_init_concentration", train_args.get("dirichlet_concentration_min", 1.0), ) ) pad_choice = str(train_args.get("conditional_pad_token", "eos") if args.conditional_pad_token == "auto" else args.conditional_pad_token) if pad_choice == "pad": pad_id = tokenizer.pad_id if tokenizer.pad_id is not None else tokenizer.eos_id loss_on_pad = False else: pad_id = tokenizer.eos_id loss_on_pad = True model = build_model_from_ckpt(ckpt, tokenizer, max_len, device, use_ema=bool(args.use_ema)) dataset = load_elf_conditional_dataset(data_path, max_records=max(0, int(args.num_samples))) collator = ELFConditionalCollator(pad_id, max_len, max_input_len, loss_on_pad=loss_on_pad) gen_len = max(0, max_len - max_input_len) out_dir = Path(args.output_dir) if args.output_dir else Path(args.checkpoint).resolve().parent / "wmt14_cond_eval" out_dir.mkdir(parents=True, exist_ok=True) print( "[eval] " f"ckpt={args.checkpoint} samples={len(dataset)} batch={args.batch_size} " f"max_len={max_len} max_input_len={max_input_len} gen_len={gen_len} " f"steps={infer_steps} solver={solver} decode_rule={decode_rule} temp={args.temperature} " f"early_temp={args.early_temp} late_temp={args.late_temp} noise={noise_mode} " f"pad={pad_choice} device={device}", flush=True, ) rows: list[dict[str, Any]] = [] t0 = time.time() for start in range(0, len(dataset), int(args.batch_size)): examples = [dataset[i] for i in range(start, min(start + int(args.batch_size), len(dataset)))] batch = collator(examples) ids = batch["ids"].to(device) cond_mask = batch["cond_seq_mask"].to(device) cond_lens = cond_mask.long().sum(dim=1) pos = torch.arange(max_len, device=device).unsqueeze(0) active_len = (cond_lens + gen_len).clamp_max(max_len) attn_mask = pos < active_len.unsqueeze(1) prefix_probs = smooth_onehot(ids, model.vocab_size, 1.0, args.eps) noise = sample_noise_simplex( (ids.size(0), max_len), model.vocab_size, device, args.eps, noise_mode=noise_mode, target_prob=target_prob, noise_sigma=noise_sigma, dirichlet_concentration=dirichlet_conc, ) init = torch.where(cond_mask.unsqueeze(-1), prefix_probs, noise) final_probs = conditional_decode( model, init, attn_mask, cond_mask, prefix_probs, infer_steps=infer_steps, damping=damping, max_gamma=max_gamma, eps=args.eps, solver=solver, decode_rule=decode_rule, model_t_mode=args.model_t_mode, temperature=args.temperature, early_temp=args.early_temp, late_temp=args.late_temp, temp_end=args.temp_end, temp_power=args.temp_power, tail_temp=args.tail_temp, support_power=args.support_power, semantic_power=args.semantic_power, hybrid_switch=args.hybrid_switch, c_min=args.c_min, c_max=args.c_max, final_endpoint_blend=args.final_endpoint_blend, pad_id=pad_id, ) final_ids = final_probs.argmax(dim=-1) for local_i, ex in enumerate(examples): gen_text = decode_target_segment( tokenizer, final_ids[local_i], int(cond_lens[local_i].item()), gen_len, detokenizer=args.detokenizer, ) row = { "id": int(ex.get("index", start + local_i)), "source": ex.get("input", ""), "reference": ex.get("target", ""), "generated": gen_text, "cond_len": int(cond_lens[local_i].item()), } rows.append(row) print(f"[eval] generated={len(rows)}/{len(dataset)}", flush=True) generations_path = out_dir / "generations.jsonl" with generations_path.open("w", encoding="utf-8") as f: for row in rows: f.write(json.dumps(row, ensure_ascii=False) + "\n") hypotheses = [row["generated"] for row in rows] references = [row["reference"] for row in rows] metrics = compute_text_metrics(hypotheses, references) metrics.update( { "checkpoint": str(args.checkpoint), "step": int(ckpt.get("step", -1)), "use_ema": bool(args.use_ema and "ema_model" in ckpt), "data_path": str(data_path), "tokenizer_path": str(tokenizer_path), "max_len": max_len, "max_input_len": max_input_len, "gen_len": gen_len, "infer_steps": infer_steps, "decode_solver": solver, "decode_rule": decode_rule, "model_t_mode": args.model_t_mode, "decode_damping": damping, "max_gamma": max_gamma, "temperature": float(args.temperature), "early_temp": float(args.early_temp), "late_temp": float(args.late_temp), "temp_end": float(args.temp_end), "temp_power": float(args.temp_power), "tail_temp": float(args.tail_temp), "support_power": float(args.support_power), "semantic_power": float(args.semantic_power), "hybrid_switch": float(args.hybrid_switch), "c_min": float(args.c_min), "c_max": float(args.c_max), "final_endpoint_blend": float(args.final_endpoint_blend), "noise_init": noise_mode, "target_prob_for_noise": target_prob, "noise_sigma": noise_sigma, "dirichlet_init_concentration": dirichlet_conc, "conditional_pad_token": pad_choice, "elapsed_sec": time.time() - t0, "generations_path": str(generations_path), } ) summary_path = out_dir / "summary.json" summary_path.write_text(json.dumps(metrics, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") for row in rows[: max(0, int(args.print_samples))]: print( "\n--- sample {id} ---\nSRC: {source}\nREF: {reference}\nGEN: {generated}".format(**row), flush=True, ) print( "[metrics] " f"BLEU={metrics.get('bleu')} chrF={metrics.get('chrf')} " f"ROUGE1={metrics.get('rouge1')} ROUGE2={metrics.get('rouge2')} ROUGEL={metrics.get('rougeL')} " f"empty={metrics.get('empty_rate'):.4f}", flush=True, ) print(f"[eval] wrote {summary_path}", flush=True) if __name__ == "__main__": main()