| |
| """Per-sample paper analysis for LatentASR. |
| |
| The script runs three paths on one ASR split: |
| 1. frozen baseline, |
| 2. LatentASR with the deployed halting threshold, |
| 3. LatentASR with forced full compute. |
| |
| It writes per-sample JSON plus a LaTeX snippet containing difficulty bins, |
| gate-quality diagnostics, qualitative examples, and latent-delta statistics. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| import math |
| import os |
| import sys |
| from pathlib import Path |
| from typing import Any, Dict, Iterable, List, Optional |
|
|
| import numpy as np |
| import torch |
| from datasets import Audio, load_dataset |
| from jiwer import cer, wer |
| from tqdm import tqdm |
|
|
| REPO_ROOT = Path(__file__).resolve().parents[2] |
| if str(REPO_ROOT) not in sys.path: |
| sys.path.insert(0, str(REPO_ROOT)) |
|
|
| from eval import ( |
| build_base_model_bundle, |
| build_latent_bundle, |
| clean_prediction, |
| configure_text_normalizer, |
| normalize_text, |
| ) |
|
|
|
|
| MODEL_ID = "Qwen/Qwen3-ASR-0.6B" |
| LATENT_CKPT = "eval_runs/paper_tbd_retrain_20260518/checkpoints/activation_500/activation_500_epoch10.pth" |
| MAX_NEW_TOKENS = 128 |
|
|
|
|
| DATASET_PRESETS = { |
| "fleurs_en_us": { |
| "dataset_name": "google/fleurs", |
| "config": "en_us", |
| "split": "test", |
| "text_columns": ["transcription", "raw_transcription", "sentence", "text"], |
| "normalizer": "english", |
| "language_hint": "English", |
| "label": r"FLEURS (\texttt{en\_us})", |
| "source_label": "FLEURS Q4", |
| }, |
| "voxpopuli_en": { |
| "dataset_name": "facebook/voxpopuli", |
| "config": "en", |
| "split": "test", |
| "text_columns": ["normalized_text", "raw_text", "text", "sentence"], |
| "normalizer": "english", |
| "language_hint": "English", |
| "label": r"VoxPopuli (\texttt{en})", |
| "source_label": "VoxPopuli Q4", |
| }, |
| } |
|
|
|
|
| def choose_device() -> str: |
| return "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
| def choose_dtype(device: str) -> torch.dtype: |
| return torch.bfloat16 if device == "cuda" and torch.cuda.is_bf16_supported() else torch.float16 |
|
|
|
|
| def tensor_to_float(value: Any) -> Optional[float]: |
| if value is None: |
| return None |
| if torch.is_tensor(value): |
| if value.numel() == 0: |
| return None |
| return float(value.detach().flatten()[0].cpu().item()) |
| try: |
| return float(value) |
| except Exception: |
| return None |
|
|
|
|
| def tensor_to_list(value: Any) -> List[float]: |
| if value is None: |
| return [] |
| if torch.is_tensor(value): |
| return [float(x) for x in value.detach().flatten().cpu().tolist()] |
| if isinstance(value, (list, tuple)): |
| return [float(x) for x in value] |
| return [] |
|
|
|
|
| def get_ref(sample: Dict[str, Any], text_columns: List[str]) -> str: |
| for col in text_columns: |
| value = sample.get(col) |
| if isinstance(value, str) and value.strip(): |
| return value |
| return "" |
|
|
|
|
| def sample_to_features(model: Any, processor: Any, sample: Dict[str, Any]) -> Optional[Dict[str, Any]]: |
| audio = sample.get("audio") |
| if not isinstance(audio, dict) or "array" not in audio or "sampling_rate" not in audio: |
| return None |
| audio_array = np.array(audio["array"], dtype=np.float64) |
| target_dtype = model.thinker.dtype if hasattr(model.thinker, "dtype") else torch.float32 |
| feat_out = processor.feature_extractor( |
| audio_array, |
| sampling_rate=audio["sampling_rate"], |
| return_attention_mask=True, |
| ) |
| device = model.base_model.device |
| feats = torch.tensor(feat_out.input_features[0], dtype=target_dtype, device=device).unsqueeze(0) |
| n_frames = feats.size(-1) |
| if getattr(feat_out, "attention_mask", None) is not None: |
| raw_mask = feat_out.attention_mask[0] |
| if not isinstance(raw_mask, (list, torch.Tensor)): |
| raw_mask = list(raw_mask) |
| if isinstance(raw_mask, torch.Tensor): |
| raw_mask = raw_mask.long() |
| else: |
| raw_mask = torch.tensor(raw_mask, dtype=torch.long) |
| if raw_mask.size(-1) < n_frames: |
| raw_mask = torch.cat([raw_mask, torch.zeros(n_frames - raw_mask.size(-1), dtype=torch.long)]) |
| else: |
| raw_mask = raw_mask[:n_frames] |
| feature_attention_mask = raw_mask.to(device=device).unsqueeze(0) |
| else: |
| feature_attention_mask = torch.ones((1, n_frames), dtype=torch.long, device=device) |
| if int(feature_attention_mask.sum().item()) < 10: |
| return None |
| return {"feats": feats, "feature_attention_mask": feature_attention_mask} |
|
|
|
|
| @torch.no_grad() |
| def transcribe( |
| model: Any, |
| processor: Any, |
| sample: Dict[str, Any], |
| *, |
| use_baseline: bool, |
| theta: float, |
| language_hint: str, |
| ) -> Optional[Dict[str, Any]]: |
| feats = sample_to_features(model, processor, sample) |
| if feats is None: |
| return None |
| prompt_text = f"Transcribe the {language_hint} audio into text." if language_hint else "Transcribe the audio into text." |
| gen_kwargs = { |
| "feature_attention_mask": feats["feature_attention_mask"], |
| "max_new_tokens": MAX_NEW_TOKENS, |
| "use_baseline": use_baseline, |
| "return_thoughts": False, |
| "return_stats": True, |
| "do_sample": False, |
| "eos_token_id": [151645, 151643], |
| "num_beams": 1, |
| "language_hint": language_hint, |
| "prompt_text": prompt_text, |
| "dynamic_halt_threshold": theta, |
| } |
| out = model.generate(feats["feats"], **gen_kwargs) |
| if isinstance(out, tuple): |
| gen_ids = out[0] |
| stats = out[1] if len(out) > 1 and isinstance(out[1], dict) else {} |
| else: |
| gen_ids = out |
| stats = {} |
| ids = gen_ids[0] |
| eos_id = processor.tokenizer.eos_token_id |
| if eos_id is not None and (ids == eos_id).any(): |
| eos_pos = (ids == eos_id).nonzero(as_tuple=True)[0][0] |
| ids = ids[:eos_pos] |
| raw = processor.tokenizer.decode(ids, skip_special_tokens=True) |
| pred = clean_prediction(raw) |
| deq = tensor_to_float(stats.get("deq_iters")) |
| return { |
| "pred": pred, |
| "pred_norm": normalize_text(pred), |
| "stats": { |
| "deq_iters": 0 if deq is None else int(round(deq)), |
| "skipped": bool(stats.get("skipped", False)) if stats else False, |
| "v_preds": tensor_to_list(stats.get("v_preds")), |
| "scaled_norm_mean": tensor_to_list(stats.get("scaled_norm_mean")), |
| "step_cos": tensor_to_float(stats.get("step_cos")), |
| "diff_norm": tensor_to_float(stats.get("diff_norm")), |
| }, |
| } |
|
|
|
|
| def load_rows(path: Path) -> List[Dict[str, Any]]: |
| if not path.exists(): |
| return [] |
| return json.loads(path.read_text(encoding="utf-8")) |
|
|
|
|
| def save_rows(path: Path, rows: List[Dict[str, Any]]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps(rows, indent=2, ensure_ascii=False), encoding="utf-8") |
|
|
|
|
| def iter_dataset(preset: Dict[str, Any], streaming: bool) -> Iterable[Dict[str, Any]]: |
| ds = load_dataset( |
| preset["dataset_name"], |
| preset["config"], |
| split=preset["split"], |
| streaming=streaming, |
| trust_remote_code=True, |
| ) |
| if not streaming: |
| ds = ds.cast_column("audio", Audio(sampling_rate=16000)) |
| return ds |
|
|
|
|
| def collect_refs(preset: Dict[str, Any], streaming: bool) -> List[Dict[str, Any]]: |
| refs: List[Dict[str, Any]] = [] |
| for idx, sample in enumerate(tqdm(iter_dataset(preset, streaming), desc="refs")): |
| ref_raw = get_ref(sample, preset["text_columns"]) |
| ref_norm = normalize_text(ref_raw) |
| if ref_norm: |
| refs.append({"idx": idx, "ref_raw": ref_raw, "ref_norm": ref_norm}) |
| return refs |
|
|
|
|
| def run_baseline(args: argparse.Namespace, preset: Dict[str, Any], refs: List[Dict[str, Any]], rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: |
| done = {r["idx"] for r in rows if "baseline_norm" in r} |
| if len(done) == len(refs): |
| return rows |
| by_idx = {r["idx"]: r for r in rows} |
| ref_idx = {r["idx"]: r for r in refs} |
| bundle = build_base_model_bundle(args.model_id, args.device, args.dtype) |
| for idx, sample in enumerate(tqdm(iter_dataset(preset, args.streaming), desc="baseline")): |
| if idx not in ref_idx or idx in done: |
| continue |
| out = transcribe( |
| bundle.model, |
| bundle.processor, |
| sample, |
| use_baseline=True, |
| theta=args.theta, |
| language_hint=preset["language_hint"], |
| ) |
| if out is None: |
| continue |
| row = by_idx.setdefault(idx, {"idx": idx, **ref_idx[idx]}) |
| row["baseline_pred"] = out["pred"] |
| row["baseline_norm"] = out["pred_norm"] |
| save_rows(args.out_json, sorted(by_idx.values(), key=lambda x: x["idx"])) |
| del bundle |
| gc.collect() |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| return sorted(by_idx.values(), key=lambda x: x["idx"]) |
|
|
|
|
| def run_latent_path( |
| args: argparse.Namespace, |
| preset: Dict[str, Any], |
| rows: List[Dict[str, Any]], |
| *, |
| theta: float, |
| pred_key: str, |
| norm_key: str, |
| stats_key: str, |
| desc: str, |
| ) -> List[Dict[str, Any]]: |
| eligible = {r["idx"] for r in rows if "baseline_norm" in r} |
| done = {r["idx"] for r in rows if norm_key in r} |
| if eligible and done == eligible: |
| return rows |
| by_idx = {r["idx"]: r for r in rows} |
| bundle = build_latent_bundle(args.model_id, args.latent_ckpt, args.n_latent, args.device, args.dtype) |
| for idx, sample in enumerate(tqdm(iter_dataset(preset, args.streaming), desc=desc)): |
| if idx not in eligible or idx in done: |
| continue |
| out = transcribe( |
| bundle.model, |
| bundle.processor, |
| sample, |
| use_baseline=False, |
| theta=theta, |
| language_hint=preset["language_hint"], |
| ) |
| if out is None: |
| continue |
| row = by_idx[idx] |
| row[pred_key] = out["pred"] |
| row[norm_key] = out["pred_norm"] |
| row[stats_key] = out["stats"] |
| save_rows(args.out_json, sorted(by_idx.values(), key=lambda x: x["idx"])) |
| del bundle |
| gc.collect() |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| return sorted(by_idx.values(), key=lambda x: x["idx"]) |
|
|
|
|
| def group_wer(group: List[Dict[str, Any]], key: str) -> float: |
| if not group: |
| return 0.0 |
| return 100 * wer([r["ref_norm"] for r in group], [r[key] for r in group]) |
|
|
|
|
| def group_cer(group: List[Dict[str, Any]], key: str) -> float: |
| if not group: |
| return 0.0 |
| return 100 * cer([r["ref_norm"] for r in group], [r[key] for r in group]) |
|
|
|
|
| def esc_latex(text: str) -> str: |
| replacements = { |
| "\\": r"\textbackslash{}", |
| "&": r"\&", |
| "%": r"\%", |
| "$": r"\$", |
| "#": r"\#", |
| "_": r"\_", |
| "{": r"\{", |
| "}": r"\}", |
| "~": r"\textasciitilde{}", |
| "^": r"\textasciicircum{}", |
| } |
| return "".join(replacements.get(ch, ch) for ch in text) |
|
|
|
|
| def format_delta(value: float, bold_negative: bool = False) -> str: |
| if bold_negative and value < 0: |
| return rf"$\boldsymbol{{{value:+.2f}}}$" |
| return rf"${value:+.2f}$" |
|
|
|
|
| def write_latex(args: argparse.Namespace, preset: Dict[str, Any], rows: List[Dict[str, Any]]) -> None: |
| final_rows = [ |
| r for r in rows |
| if all(k in r for k in ("baseline_norm", "latent_norm", "full_norm")) |
| ] |
| for r in final_rows: |
| ref = r["ref_norm"] |
| r["baseline_wer"] = float(wer(ref, r["baseline_norm"])) |
| r["latent_wer"] = float(wer(ref, r["latent_norm"])) |
| r["full_wer"] = float(wer(ref, r["full_norm"])) |
| r["baseline_cer"] = float(cer(ref, r["baseline_norm"])) |
| r["latent_cer"] = float(cer(ref, r["latent_norm"])) |
| r["full_cer"] = float(cer(ref, r["full_norm"])) |
|
|
| base_wer = group_wer(final_rows, "baseline_norm") |
| lat_wer = group_wer(final_rows, "latent_norm") |
| full_wer = group_wer(final_rows, "full_norm") |
| base_cer = group_cer(final_rows, "baseline_norm") |
| lat_cer = group_cer(final_rows, "latent_norm") |
| full_cer = group_cer(final_rows, "full_norm") |
|
|
| sorted_rows = sorted(final_rows, key=lambda r: (r["baseline_wer"], r["idx"])) |
| n = len(sorted_rows) |
| bins = [] |
| for q in range(4): |
| part = sorted_rows[math.floor(q * n / 4): math.floor((q + 1) * n / 4)] |
| bw = group_wer(part, "baseline_norm") |
| lw = group_wer(part, "latent_norm") |
| skip_q = 100 * sum(1 for r in part if r["latent_stats"].get("deq_iters", 0) == 0) / max(1, len(part)) |
| bins.append((q + 1, len(part), bw, lw, lw - bw, skip_q)) |
|
|
| skipped = [r for r in final_rows if r["latent_stats"].get("deq_iters", 0) == 0] |
| processed = [r for r in final_rows if r["latent_stats"].get("deq_iters", 0) > 0] |
|
|
| skip_base = group_wer(skipped, "baseline_norm") |
| skip_full = group_wer(skipped, "full_norm") |
| proc_base = group_wer(processed, "baseline_norm") |
| proc_lat = group_wer(processed, "latent_norm") |
| proc_full = group_wer(processed, "full_norm") |
|
|
| step_counts = {i: 0 for i in range(args.n_latent + 1)} |
| for r in final_rows: |
| deq = int(r["latent_stats"].get("deq_iters", 0)) |
| step_counts[deq] = step_counts.get(deq, 0) + 1 |
| step_rates = {k: 100 * v / max(1, len(final_rows)) for k, v in step_counts.items()} |
| avg_steps = sum(k * v for k, v in step_counts.items()) / max(1, len(final_rows)) |
|
|
| processed_full = processed |
| norms_by_step: List[List[float]] = [[] for _ in range(args.n_latent)] |
| cos_vals: List[float] = [] |
| diff_vals: List[float] = [] |
| for r in processed_full: |
| norms = r["full_stats"].get("scaled_norm_mean", []) |
| for i, val in enumerate(norms[:args.n_latent]): |
| norms_by_step[i].append(float(val)) |
| if r["full_stats"].get("step_cos") is not None: |
| cos_vals.append(float(r["full_stats"]["step_cos"])) |
| if r["full_stats"].get("diff_norm") is not None: |
| diff_vals.append(float(r["full_stats"]["diff_norm"])) |
| step_scale_means = [float(np.mean(vals)) if vals else 0.0 for vals in norms_by_step] |
| step_scale_text = ", ".join(f"${v:.4f}$" for v in step_scale_means) |
|
|
| def dist_stats(vals: List[float]) -> Dict[str, float]: |
| if not vals: |
| return { |
| "mean": 0.0, |
| "std": 0.0, |
| "p25": 0.0, |
| "median": 0.0, |
| "p75": 0.0, |
| "min": 0.0, |
| "max": 0.0, |
| } |
| arr = np.asarray(vals, dtype=np.float64) |
| return { |
| "mean": float(np.mean(arr)), |
| "std": float(np.std(arr)), |
| "p25": float(np.percentile(arr, 25)), |
| "median": float(np.percentile(arr, 50)), |
| "p75": float(np.percentile(arr, 75)), |
| "min": float(np.min(arr)), |
| "max": float(np.max(arr)), |
| } |
|
|
| cos_stat = dist_stats(cos_vals) |
| diff_stat = dist_stats(diff_vals) |
|
|
| examples = [ |
| r for r in final_rows |
| if r["baseline_wer"] > r["latent_wer"] |
| and r["latent_stats"].get("deq_iters", 0) >= 1 |
| and len(r["ref_norm"].split()) >= 6 |
| ] |
| examples.sort(key=lambda r: (r["baseline_wer"] - r["latent_wer"], r["baseline_wer"]), reverse=True) |
| examples = examples[:4] |
|
|
| lines: List[str] = [] |
| lines.append("% ================================================================\n") |
| lines.append(f"% Per-sample analysis generated for {args.dataset_key}\n") |
| lines.append(f"% Samples used: {len(final_rows)}\n") |
| lines.append("% ================================================================\n\n") |
| lines.append("\\subsection{Analysis}\n") |
| lines.append("\\label{sec:analysis}\n\n") |
| lines.append("\\textbf{Difficulty-Binned Reductions.}\\quad\n") |
| lines.append( |
| f"We partition the {preset['label']} test set into four equal-sized bins by per-utterance " |
| f"Baseline WER and recompute WER within each bin (Table~\\ref{{tab:difficulty_bins}}). " |
| f"The aggregate result is small but positive: \\method{{}} reduces WER from " |
| f"${base_wer:.2f}\\%$ to ${lat_wer:.2f}\\%$ at $\\theta{{=}}{args.theta:.1f}$, " |
| f"while forced full compute reaches ${full_wer:.2f}\\%$. " |
| ) |
| best_bin = min(bins, key=lambda x: x[4]) |
| lines.append( |
| f"The largest reduction appears in Q{best_bin[0]}, where $\\Delta$WER is " |
| f"${best_bin[4]:+.2f}$~pp. This confirms that the average gain should not be read " |
| "as a uniform per-utterance improvement; latent scaling mainly changes the subset " |
| "where the frozen baseline leaves residual errors.\n\n" |
| ) |
| lines.append("\\begin{table}[ht]\n") |
| lines.append( |
| f" \\caption{{Difficulty-binned analysis on {preset['label']} at $\\theta{{=}}{args.theta:.1f}$ " |
| f"({len(final_rows):,} utterances total). Utterances are partitioned into Baseline-WER quartiles. " |
| "$\\Delta$WER denotes \\method{} minus Baseline, so negative values indicate improvement.}\n" |
| ) |
| lines.append(" \\label{tab:difficulty_bins}\n") |
| lines.append(" \\centering\n") |
| lines.append(" \\resizebox{\\columnwidth}{!}{\n") |
| lines.append(" \\begin{tabular}{l c c c c c}\n") |
| lines.append(" \\toprule\n") |
| lines.append(" \\textbf{Bin} & \\textbf{\\#utts} & \\textbf{Baseline WER (\\%)} & \\textbf{\\method{} WER (\\%)} & \\textbf{$\\Delta$WER (pp)} & \\textbf{Skip (\\%)} \\\\\n") |
| lines.append(" \\midrule\n") |
| labels = ["Q1 (easiest)", "Q2", "Q3", "Q4 (hardest)"] |
| for q, count, bw, lw, d, skip_q in bins: |
| method_cell = rf"\textbf{{{lw:.2f}}}" if d < 0 else f"{lw:.2f}" |
| lines.append( |
| f" {labels[q - 1]} & {count:,} & {bw:.2f} & {method_cell} & " |
| f"{format_delta(d, bold_negative=True)} & {skip_q:.1f} \\\\\n" |
| ) |
| lines.append(" \\bottomrule\n") |
| lines.append(" \\end{tabular}\n") |
| lines.append(" }\n") |
| lines.append("\\end{table}\n\n") |
|
|
| lines.append("\\textbf{Value Head Decision Quality.}\\quad\n") |
| lines.append( |
| "The step distribution shows how much compute the Value Head allocates: " |
| f"at $\\theta{{=}}{args.theta:.1f}$, it skips {step_rates.get(0, 0.0):.1f}\\% of utterances " |
| f"and uses an average of {avg_steps:.2f} latent steps. " |
| "We further test selectivity by forcing the full $N{=}4$ path on the utterances " |
| "that the deployed gate skips. Table~\\ref{tab:gate_quality} reports the actual " |
| "deployed change and this counterfactual full-compute change.\n\n" |
| ) |
| lines.append("\\begin{table}[ht]\n") |
| lines.append( |
| f" \\caption{{Value Head decision quality on {preset['label']} at $\\theta{{=}}{args.theta:.1f}$. " |
| "\\textbf{Counterfactual $\\Delta$WER} forces the $N{=}4$ latent path on each subset.}\n" |
| ) |
| lines.append(" \\label{tab:gate_quality}\n") |
| lines.append(" \\centering\n") |
| lines.append(" \\resizebox{\\columnwidth}{!}{\n") |
| lines.append(" \\begin{tabular}{l c c c}\n") |
| lines.append(" \\toprule\n") |
| lines.append(" \\textbf{Subset (at $\\theta{=}0.0$)} & \\textbf{\\#utts} & \\textbf{Actual $\\Delta$WER (pp)} & \\textbf{Counterfactual $\\Delta$WER (pp)} \\\\\n") |
| lines.append(" \\midrule\n") |
| lines.append(f" Skipped ($v_0 < 0$) & {len(skipped):,} & $0.00$ (by construction) & {format_delta(skip_full - skip_base, True)} \\\\\n") |
| lines.append(f" Processed ($v_0 \\geq 0$) & {len(processed):,} & {format_delta(proc_lat - proc_base, True)} & {format_delta(proc_full - proc_base, True)} \\\\\n") |
| lines.append(" \\bottomrule\n") |
| lines.append(" \\end{tabular}\n") |
| lines.append(" }\n") |
| lines.append("\\end{table}\n\n") |
|
|
| lines.append("\\textbf{Step Allocation.}\\quad\n") |
| lines.append( |
| f"Table~\\ref{{tab:vox_step_dist}} gives the full $N$-step distribution on {preset['label']}. " |
| "Compared with forced full compute, the deployed policy keeps most examples away from " |
| "the deepest path while retaining the aggregate WER reduction.\n\n" |
| ) |
| lines.append("\\begin{table}[ht]\n") |
| lines.append( |
| f" \\caption{{N-step distribution on {preset['label']} at $\\theta{{=}}{args.theta:.1f}$.}}\n" |
| ) |
| lines.append(" \\label{tab:vox_step_dist}\n") |
| lines.append(" \\centering\n") |
| lines.append(" \\resizebox{0.9\\columnwidth}{!}{\n") |
| lines.append(" \\begin{tabular}{l c c c c c c}\n") |
| lines.append(" \\toprule\n") |
| lines.append(" \\textbf{Dataset} & \\textbf{Avg. steps} & \\textbf{N=0} & \\textbf{N=1} & \\textbf{N=2} & \\textbf{N=3} & \\textbf{N=4} \\\\\n") |
| lines.append(" \\midrule\n") |
| lines.append( |
| f" {preset['label']} & {avg_steps:.2f} & " |
| f"{step_rates.get(0, 0.0):.1f}\\% & {step_rates.get(1, 0.0):.1f}\\% & " |
| f"{step_rates.get(2, 0.0):.1f}\\% & {step_rates.get(3, 0.0):.1f}\\% & " |
| f"{step_rates.get(4, 0.0):.1f}\\% \\\\\n" |
| ) |
| lines.append(" \\bottomrule\n") |
| lines.append(" \\end{tabular}\n") |
| lines.append(" }\n") |
| lines.append("\\end{table}\n\n") |
|
|
| if examples: |
| lines.append("\\textbf{Qualitative Examples.}\\quad\n") |
| lines.append( |
| f"Table~\\ref{{tab:qualitative}} shows hard-bin {preset['label']} utterances where " |
| "the latent loop changes the transcript.\n\n" |
| ) |
| lines.append("\\begin{table}[ht]\n") |
| lines.append( |
| f" \\caption{{Qualitative examples from {preset['label']} hard bins.}}\n" |
| ) |
| lines.append(" \\label{tab:qualitative}\n") |
| lines.append(" \\centering\n") |
| lines.append(" \\resizebox{\\columnwidth}{!}{\n") |
| lines.append(" \\begin{tabular}{p{0.13\\columnwidth} p{0.27\\columnwidth} p{0.27\\columnwidth} p{0.27\\columnwidth}}\n") |
| lines.append(" \\toprule\n") |
| lines.append(" \\textbf{Source} & \\textbf{Reference} & \\textbf{Baseline} & \\textbf{\\method{}} \\\\\n") |
| lines.append(" \\midrule\n") |
| for r in examples: |
| lines.append( |
| f" {preset['source_label']} & {esc_latex(r['ref_raw'])} & " |
| f"{esc_latex(r['baseline_pred'])} & {esc_latex(r['latent_pred'])} \\\\\n" |
| ) |
| lines.append(" \\bottomrule\n") |
| lines.append(" \\end{tabular}\n") |
| lines.append(" }\n") |
| lines.append("\\end{table}\n\n") |
|
|
| lines.append("\\textbf{Refinement-Path Diagnostics.}\\quad\n") |
| lines.append( |
| "We recompute the forced-full $N{=}4$ path on the " |
| f"{preset['label']} processed subset ({len(processed_full):,} utterances with $N{{>}}0$ " |
| f"under $\\theta{{=}}{args.theta:.1f}$). The scaled delta norms are identical across " |
| "utterances because each delta is $L_2$-normalized and multiplied by the learned " |
| "per-step scale; they therefore measure the bounded step-size constraint rather " |
| f"than dataset-specific refinement behavior. For this run, the per-step scales are {step_scale_text}. " |
| "To characterize sample-dependent behavior, Table~\\ref{tab:refinement_dynamics} " |
| "instead reports the distribution of consecutive-delta cosine and consecutive-delta " |
| "difference norm. " |
| f"The cosine range is ${cos_stat['min']:.4f}$--${cos_stat['max']:.4f}$ and the " |
| f"difference-norm range is ${diff_stat['min']:.4f}$--${diff_stat['max']:.4f}$, " |
| "confirming that the forced refinement path is not a constant copied trajectory " |
| "while the update magnitudes remain bounded.\n\n" |
| ) |
| lines.append("\\begin{table}[ht]\n") |
| lines.append( |
| f" \\caption{{Forced-full refinement diagnostics on the {preset['label']} processed subset " |
| f"({len(processed_full):,} utterances with $N{{>}}0$ under $\\theta{{=}}{args.theta:.1f}$). " |
| "Scaled delta norms are fixed by the learned step scales; cosine and difference " |
| "statistics vary across utterances.}\n" |
| ) |
| lines.append(" \\label{tab:refinement_dynamics}\n") |
| lines.append(" \\centering\n") |
| lines.append(" \\resizebox{\\columnwidth}{!}{\n") |
| lines.append(" \\begin{tabular}{l c c c c c}\n") |
| lines.append(" \\toprule\n") |
| lines.append(r" \textbf{Metric} & \textbf{Mean} & \textbf{Std.} & \textbf{P25} & \textbf{Median} & \textbf{P75} \\" + "\n") |
| lines.append(" \\midrule\n") |
| lines.append( |
| f" Consecutive-delta cosine & {cos_stat['mean']:.4f} & {cos_stat['std']:.4f} & " |
| f"{cos_stat['p25']:.4f} & {cos_stat['median']:.4f} & {cos_stat['p75']:.4f} " |
| + r"\\" |
| + "\n" |
| ) |
| lines.append( |
| f" Consecutive-delta diff. norm & {diff_stat['mean']:.4f} & {diff_stat['std']:.4f} & " |
| f"{diff_stat['p25']:.4f} & {diff_stat['median']:.4f} & {diff_stat['p75']:.4f} " |
| + r"\\" |
| + "\n" |
| ) |
| lines.append(" \\bottomrule\n") |
| lines.append(" \\end{tabular}\n") |
| lines.append(" }\n") |
| lines.append("\\end{table}\n\n") |
|
|
| lines.append("% Overall metrics for cross-checking:\n") |
| lines.append(f"% Baseline WER/CER: {base_wer:.4f}/{base_cer:.4f}\n") |
| lines.append(f"% Latent theta={args.theta:.1f} WER/CER: {lat_wer:.4f}/{lat_cer:.4f}\n") |
| lines.append(f"% Forced full WER/CER: {full_wer:.4f}/{full_cer:.4f}\n") |
|
|
| args.out_tex.write_text("".join(lines), encoding="utf-8") |
| save_rows(args.out_json, final_rows) |
| print(f"wrote {args.out_json}") |
| print(f"wrote {args.out_tex}") |
| print(f"samples={len(final_rows)} baseline_wer={base_wer:.4f} latent_wer={lat_wer:.4f} full_wer={full_wer:.4f}") |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--dataset-key", choices=sorted(DATASET_PRESETS), default="voxpopuli_en") |
| parser.add_argument("--model-id", default=MODEL_ID) |
| parser.add_argument("--latent-ckpt", default=LATENT_CKPT) |
| parser.add_argument("--n-latent", type=int, default=4) |
| parser.add_argument("--theta", type=float, default=0.0) |
| parser.add_argument("--full-theta", type=float, default=-2.0) |
| parser.add_argument("--streaming", action=argparse.BooleanOptionalAction, default=True) |
| parser.add_argument("--out-dir", type=Path, default=Path("eval_runs/paper_activation500_voxpopuli_analysis")) |
| args = parser.parse_args() |
| args.out_dir.mkdir(parents=True, exist_ok=True) |
| args.out_json = args.out_dir / f"{args.dataset_key}_per_sample.json" |
| args.out_tex = args.out_dir / f"{args.dataset_key}_analysis_latex.tex" |
| args.device = choose_device() |
| args.dtype = choose_dtype(args.device) |
| return args |
|
|
|
|
| def main() -> None: |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") |
| args = parse_args() |
| preset = DATASET_PRESETS[args.dataset_key] |
| configure_text_normalizer(preset["normalizer"]) |
|
|
| rows = load_rows(args.out_json) |
| refs = collect_refs(preset, args.streaming) |
| existing = {r["idx"]: r for r in rows} |
| for ref in refs: |
| existing.setdefault(ref["idx"], {"idx": ref["idx"], **ref}) |
| rows = sorted(existing.values(), key=lambda x: x["idx"]) |
| save_rows(args.out_json, rows) |
|
|
| rows = run_baseline(args, preset, refs, rows) |
| rows = run_latent_path( |
| args, |
| preset, |
| rows, |
| theta=args.theta, |
| pred_key="latent_pred", |
| norm_key="latent_norm", |
| stats_key="latent_stats", |
| desc=f"latent_theta{args.theta:g}", |
| ) |
| rows = run_latent_path( |
| args, |
| preset, |
| rows, |
| theta=args.full_theta, |
| pred_key="full_pred", |
| norm_key="full_norm", |
| stats_key="full_stats", |
| desc="latent_full", |
| ) |
| write_latex(args, preset, rows) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
