| |
| 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 <checkpoint_dir>/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: |
| 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: |
| 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() |
|
|
