""" Evaluate latent-reasoning and baseline/PEFT ASR models. Supported datasets: - SpeechTest/common_voice_16_0 - SpeechTest/peoples_speech - SpeechTest/librispeech_asr - SpeechTest/fleurs - SpeechTest/voxpopuli - SpeechTest/gigaspeech - SpeechTest/ASCEND - SpeechTest/extreme_asr_pony - openslr/librispeech_asr - google/fleurs - facebook/voxpopuli - PolyAI/minds14 - LIUM/tedlium - edinburghcstr/ami Unknown HuggingFace datasets are also accepted with generic ASR defaults: an `audio` column plus one of the fallback text columns. Example: python eval.py \ --baseline-ckpt baseline_qwen_asr_best.pth \ --prompt-tuning-ckpt prompt_tuning_qwen_asr_best \ --lora-r16-ckpt lora_r16_qwen_asr_best \ --latent-ckpt latent_qwen_asr_best.pth \ --max-samples-per-config 500 \ --output-json cv16_eval.json """ import argparse import gc import itertools import json import os import re import unicodedata from dataclasses import dataclass from typing import Dict, List, Optional, Tuple import numpy as np import torch from datasets import Audio, get_dataset_config_names, load_dataset from jiwer import cer, wer from tqdm import tqdm from qwen_asr import Qwen3ASRModel # type: ignore from model import LatentQwenASR from utils import set_seed COMMONVOICE_12_LANGS = [ "en", ] FLEURS_12_LANGS = [ "en", ] VOXPOPULI_30_LANGS = [ "en", ] QWEN_FLEURS_CORE12_LANGS = [ "en", "zh", "yue", "ar", "de", "es", "fr", "it", "ja", "ko", "pt", "ru", ] QWEN_FLEURS_EXTRA8_LANGS = [ "hi", "id", "ms", "nl", "pl", "th", "tr", "vi", ] QWEN_FLEURS_EXTRA10_LANGS = [ "cs", "da", "el", "fa", "fi", "fil", "hu", "mk", "ro", "sv", ] QWEN_FLEURS_30_LANGS = ( QWEN_FLEURS_CORE12_LANGS + QWEN_FLEURS_EXTRA8_LANGS + QWEN_FLEURS_EXTRA10_LANGS ) MLS_PUBLIC7_LANGS = ["de", "nl", "es", "fr", "it", "pl", "pt"] COMMONVOICE_CONFIG_ALIASES = { "zh": "zh-CN", "zh-cn": "zh-CN", "zh_cn": "zh-CN", "zh_tw": "zh-TW", "zh-tw": "zh-TW", } COMMONVOICE_LANGUAGE_HINTS = { "en": "English", } GENERAL_LANGUAGE_HINTS = { "ar": "Arabic", "ar_eg": "Arabic", "ar-eg": "Arabic", "ara": "Arabic", "arb": "Arabic", "cs": "Czech", "cs_cz": "Czech", "cs-cz": "Czech", "czech": "Czech", "da": "Danish", "da_dk": "Danish", "da-dk": "Danish", "danish": "Danish", "en": "English", "en_us": "English", "en-us": "English", "en_gb": "English", "en-gb": "English", "en_au": "English", "en-au": "English", "english": "English", "fr": "French", "fr_fr": "French", "fr-fr": "French", "french": "French", "de": "German", "de_de": "German", "de-de": "German", "german": "German", "es": "Spanish", "es_es": "Spanish", "es-es": "Spanish", "es_419": "Spanish", "spanish": "Spanish", "it": "Italian", "it_it": "Italian", "it-it": "Italian", "italian": "Italian", "pt": "Portuguese", "pt_pt": "Portuguese", "pt-pt": "Portuguese", "pt_br": "Portuguese", "pt-br": "Portuguese", "portuguese": "Portuguese", "nl": "Dutch", "nl_nl": "Dutch", "nl-nl": "Dutch", "dutch": "Dutch", "pl": "Polish", "pl_pl": "Polish", "pl-pl": "Polish", "polish": "Polish", "zh": "Chinese", "zh_cn": "Chinese", "zh-cn": "Chinese", "cmn": "Chinese", "cmn_hans_cn": "Chinese", "cmn-hans-cn": "Chinese", "yue": "Cantonese", "yue_hant_hk": "Cantonese", "yue-hant-hk": "Cantonese", "hi": "Hindi", "hi_in": "Hindi", "hi-in": "Hindi", "id": "Indonesian", "id_id": "Indonesian", "id-id": "Indonesian", "ja": "Japanese", "ja_jp": "Japanese", "ja-jp": "Japanese", "ko": "Korean", "ko_kr": "Korean", "ko-kr": "Korean", "ms": "Malay", "ms_my": "Malay", "ms-my": "Malay", "ru": "Russian", "ru_ru": "Russian", "ru-ru": "Russian", "th": "Thai", "th_th": "Thai", "th-th": "Thai", "tr": "Turkish", "tr_tr": "Turkish", "tr-tr": "Turkish", "vi": "Vietnamese", "vi_vn": "Vietnamese", "vi-vn": "Vietnamese", "el": "Greek", "el_gr": "Greek", "el-gr": "Greek", "fa": "Persian", "fa_ir": "Persian", "fa-ir": "Persian", "fi": "Finnish", "fi_fi": "Finnish", "fi-fi": "Finnish", "fil": "Filipino", "fil_ph": "Filipino", "fil-ph": "Filipino", "hu": "Hungarian", "hu_hu": "Hungarian", "hu-hu": "Hungarian", "mk": "Macedonian", "mk_mk": "Macedonian", "mk-mk": "Macedonian", "ro": "Romanian", "ro_ro": "Romanian", "ro-ro": "Romanian", "sv": "Swedish", "sv_se": "Swedish", "sv-se": "Swedish", } FALLBACK_TEXT_COLUMNS = [ "sentence", "text", "normalized_text", "raw_text", "transcript", "transcription", "english_transcription", ] # Language-code family aliases used when dataset config IDs differ from # requested language IDs (e.g., zh -> cmn_hans_cn, fil -> tl, ms -> msa). LANGUAGE_CODE_ALIASES = { "en": ["eng"], "zh": ["zho", "cmn", "cmn_hans_cn", "cmn_hans"], "zh_tw": ["cmn_hant_tw", "cmn_hant", "zho_hant", "zh_hant_tw"], "yue": ["yue_hant_hk", "zh_hk", "zh_yue"], "ar": ["ara", "arb"], "de": ["deu", "ger"], "es": ["spa", "es_419"], "fr": ["fra", "fre"], "it": ["ita"], "ja": ["jpn"], "ko": ["kor"], "pt": ["por"], "ru": ["rus"], "id": ["ind", "in"], "ms": ["msa", "may", "zsm"], "fil": ["tgl", "tl"], "th": ["tha"], "vi": ["vie"], "tr": ["tur"], "hi": ["hin"], "nl": ["nld", "dut"], "sv": ["swe"], "da": ["dan"], "fi": ["fin"], "pl": ["pol"], "cs": ["ces", "cze"], "fa": ["fas", "per", "pes"], "el": ["ell", "gre"], "hu": ["hun"], "mk": ["mkd", "mac"], "ro": ["ron", "rum"], } DATASET_PRESETS = { "SpeechTest/common_voice_16_0": { "default_strategy": "fixed_languages", "default_languages": COMMONVOICE_12_LANGS, "config_aliases": COMMONVOICE_CONFIG_ALIASES, "language_hints": COMMONVOICE_LANGUAGE_HINTS, "preferred_text_columns": ["sentence"], }, "SpeechTest/peoples_speech": { "default_strategy": "all_available", "config_aliases": {}, "language_hints": {"test": "English", "en": "English"}, "preferred_text_columns": ["text"], }, "SpeechTest/librispeech_asr": { "default_strategy": "all_available", "config_aliases": {}, "language_hints": { "clean": "English", "other": "English", "en": "English", }, "preferred_text_columns": ["text"], }, "SpeechTest/fleurs": { "default_strategy": "fixed_languages", "default_languages": FLEURS_12_LANGS, "config_aliases": {}, "language_hints": COMMONVOICE_LANGUAGE_HINTS, "preferred_text_columns": ["transcription", "raw_transcription", "text"], }, "SpeechTest/voxpopuli": { "default_strategy": "fixed_languages", "default_languages": VOXPOPULI_30_LANGS, "config_aliases": {}, "language_hints": { "en": "English", }, "preferred_text_columns": ["normalized_text", "raw_text", "text"], }, "SpeechTest/gigaspeech": { "default_strategy": "all_available", "config_aliases": {}, "language_hints": {"test": "English", "en": "English"}, "preferred_text_columns": ["text"], }, "SpeechTest/ASCEND": { "default_strategy": "all_available", "config_aliases": {}, "language_hints": {}, "preferred_text_columns": ["transcription", "text"], }, "SpeechTest/extreme_asr_pony": { "default_strategy": "all_available", "config_aliases": {}, "language_hints": {}, "preferred_text_columns": ["text"], }, "openslr/librispeech_asr": { "default_strategy": "default_configs", "default_configs": ["clean", "other"], "config_aliases": {}, "language_hints": GENERAL_LANGUAGE_HINTS, "preferred_text_columns": ["text"], }, "google/fleurs": { "default_strategy": "fixed_languages", "default_languages": FLEURS_12_LANGS, "config_aliases": {}, "language_hints": GENERAL_LANGUAGE_HINTS, "preferred_text_columns": ["transcription", "raw_transcription", "text"], }, "facebook/voxpopuli": { "default_strategy": "fixed_languages", "default_languages": VOXPOPULI_30_LANGS, "config_aliases": {}, "language_hints": GENERAL_LANGUAGE_HINTS, "preferred_text_columns": ["normalized_text", "raw_text", "text"], }, "PolyAI/minds14": { "default_strategy": "default_configs", "default_configs": ["en-US"], "config_aliases": { "en": "en-US", "en_us": "en-US", "en-us": "en-US", }, "language_hints": GENERAL_LANGUAGE_HINTS, "preferred_text_columns": ["transcription", "english_transcription", "text"], }, "facebook/multilingual_librispeech": { "default_strategy": "fixed_languages", "default_languages": MLS_PUBLIC7_LANGS, "config_aliases": { "de": "german", "nl": "dutch", "es": "spanish", "fr": "french", "it": "italian", "pl": "polish", "pt": "portuguese", }, "language_hints": GENERAL_LANGUAGE_HINTS, "preferred_text_columns": ["transcript", "text"], }, "LIUM/tedlium": { "default_strategy": "default_configs", "default_configs": ["release1"], "config_aliases": { "r1": "release1", "r2": "release2", "r3": "release3", }, "language_hints": GENERAL_LANGUAGE_HINTS, "preferred_text_columns": ["text"], }, "TwinkStart/tedlium": { "default_strategy": "default_configs", "default_configs": ["release1"], "config_aliases": { "r1": "release1", "r2": "release2", "r3": "release3", }, "language_hints": GENERAL_LANGUAGE_HINTS, "preferred_text_columns": ["text"], }, "edinburghcstr/ami": { "default_strategy": "default_configs", "default_configs": ["ihm"], "config_aliases": {}, "language_hints": GENERAL_LANGUAGE_HINTS, "preferred_text_columns": ["text", "transcript", "sentence"], }, } @dataclass class ModelBundle: model: LatentQwenASR processor: any checkpoint: Dict[str, any] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( "Evaluate five ASR settings on SpeechTest or HuggingFace ASR datasets: " "base_model, baseline_ft, prompt_tuning, lora_r16, latent_reasoning." ) ) parser.add_argument( "--model-id", type=str, default=os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-0.6B"), help="Base Qwen3-ASR model ID.", ) parser.add_argument( "--dataset-name", type=str, default="SpeechTest/common_voice_16_0", help="HuggingFace dataset name.", ) parser.add_argument( "--split", type=str, default="test", help="Dataset split to evaluate.", ) parser.add_argument( "--configs", type=str, default="", help=( "Comma-separated config names (dataset-specific). " "Default: dataset preset " "(fixed-language datasets use preset languages; " "others use preset configs or all subsets). " "Ignored when --all-configs is set." ), ) parser.add_argument( "--all-configs", action="store_true", help="Evaluate every config in the dataset.", ) parser.add_argument( "--max-configs", type=int, default=0, help="Limit number of configs (0 = no limit).", ) parser.add_argument( "--max-samples-per-config", type=int, default=0, help="Maximum test samples per config (0 = all).", ) parser.add_argument( "--streaming", action="store_true", help=( "Load datasets with HuggingFace streaming=True and iterate samples " "sequentially without materializing the split on disk." ), ) parser.add_argument( "--text-normalizer", type=str, default="english", choices=["english", "basic", "simple"], help=( "Text normalization before WER/CER. Use 'basic' or 'simple' for " "multilingual evaluations to avoid English-specific rewriting." ), ) parser.add_argument( "--max-new-tokens", type=int, default=128, help="Max new tokens for generation.", ) parser.add_argument( "--baseline-ckpt", type=str, default="", help="Path to baseline fine-tuned checkpoint (.pth).", ) parser.add_argument( "--prompt-tuning-ckpt", type=str, default="", help=( "Path to prompt-tuning checkpoint. " "Supports .pth (wrapper state_dict) or a PEFT adapter directory." ), ) parser.add_argument( "--lora-r16-ckpt", type=str, default="", help=( "Path to LoRA(rank=16) checkpoint. " "Supports .pth (wrapper state_dict) or a PEFT adapter directory." ), ) parser.add_argument( "--latent-ckpt", type=str, default="", help="Path to latent reasoning checkpoint (.pth).", ) parser.add_argument( "--skip-base-model", action="store_true", help="Skip raw pretrained model inference baseline.", ) parser.add_argument( "--skip-baseline-ft", action="store_true", help="Skip baseline fine-tuned model evaluation.", ) parser.add_argument( "--skip-prompt-tuning", action="store_true", help="Skip prompt tuning model evaluation.", ) parser.add_argument( "--skip-lora-r16", action="store_true", help="Skip LoRA r16 model evaluation.", ) parser.add_argument( "--n-latent", type=int, default=-1, help="Override latent token count. If < 0, use checkpoint n_latent.", ) parser.add_argument( "--num-beams", type=int, default=1, help="Number of beams for beam search.", ) parser.add_argument( "--dynamic-halt-threshold", type=float, default=0.0, help="Value head threshold for early halting of latent thoughts.", ) parser.add_argument( "--device", type=str, default="auto", choices=["auto", "cuda", "cpu"], help="Device to run inference.", ) parser.add_argument( "--dtype", type=str, default="auto", choices=["auto", "float32", "float16", "bfloat16"], help="Inference dtype.", ) parser.add_argument( "--seed", type=int, default=42, help="Random seed.", ) parser.add_argument( "--print-samples", type=int, default=2, help="Number of qualitative samples to print per config.", ) parser.add_argument( "--print-gate-skips", action="store_true", help=( "Also print latent_reasoning samples skipped by the value-head gate. " "Default is off to keep large sweep logs readable." ), ) parser.add_argument( "--output-json", type=str, default="", help="If set, write detailed results to this JSON file.", ) parser.add_argument( "--no-language-hint", action="store_true", help="Disable language-specific prompting and use generic ASR prompt.", ) parser.add_argument( "--prompt-template", type=str, default="Transcribe the {language} audio into text.", help=( "User prompt template. Use {language} placeholder. " "Ignored when --no-language-hint is set." ), ) parser.add_argument( "--snr-db", type=float, nargs="+", default=None, help=( "Add Gaussian noise to audio at specified SNR levels (dB). " "Can specify multiple values, e.g. --snr-db 20 10 5. " "Each SNR level will be evaluated as a separate row. " "If not set, no noise is added." ), ) return parser.parse_args() def choose_device(device_arg: str) -> str: if device_arg == "auto": return "cuda" if torch.cuda.is_available() else "cpu" if device_arg == "cuda" and not torch.cuda.is_available(): raise RuntimeError("Requested --device cuda but CUDA is not available.") return device_arg def choose_dtype(dtype_arg: str, device: str) -> torch.dtype: if dtype_arg == "float32": return torch.float32 if dtype_arg == "float16": return torch.float16 if device == "cuda" else torch.float32 if dtype_arg == "bfloat16": return torch.bfloat16 if device == "cuda" else torch.float32 # auto if device != "cuda": return torch.float32 major = torch.cuda.get_device_capability(0)[0] if major >= 8: return torch.bfloat16 return torch.float16 def get_dataset_preset(dataset_name: str) -> Dict[str, any]: return DATASET_PRESETS.get(dataset_name, {}) def resolve_dataset_name(dataset_name: str) -> str: key = (dataset_name or "").strip() if not key: raise ValueError("dataset name is empty.") if key in DATASET_PRESETS: return key lowered = key.lower() for canonical in DATASET_PRESETS: if canonical.lower() == lowered: return canonical print( f"[warn] dataset '{dataset_name}' has no local preset; " "using generic HuggingFace ASR defaults." ) return key def _map_config_alias(name: str, alias_map: Dict[str, str]) -> str: n = name.strip() if not n: return n lowered = n.lower() if lowered in alias_map: return alias_map[lowered] return n def resolve_language_hint(dataset_name: str, config_name: str) -> Optional[str]: preset = get_dataset_preset(dataset_name) hints_raw = preset.get("language_hints", {}) hints = {k.lower(): v for k, v in hints_raw.items()} lowered = config_name.lower() if lowered in hints: return hints[lowered] # Generic base-language fallback (e.g. "fr_ca" -> "fr"). base = lowered.split("-")[0].split("_")[0] for key, val in hints.items(): if key.split("-")[0].split("_")[0] == base: return val return None def _norm_cfg_key(v: str) -> str: return (v or "").strip().lower().replace("-", "_") def _base_lang(v: str) -> str: return _norm_cfg_key(v).split("_")[0] def _language_candidate_forms(lang_code: str) -> List[str]: code = _norm_cfg_key(lang_code) out = { code, code.replace("_", "-"), } base = _base_lang(code) out.add(base) # Common alias expansions for multilingual dataset configs. if code == "zh": out.update({"zh_cn", "zh-cn", "zh_hans", "zh-hans"}) if code in {"zh_tw", "zhtw"}: out.update({"zh_tw", "zh-tw", "zh_hant", "zh-hant", "zh_hk", "zh-hk"}) if code == "yue": out.update({"yue_hk", "yue-hk", "yue_hant_hk", "yue-hant-hk", "zh_yue", "zh-yue", "zh_hk", "zh-hk"}) if code == "fil": out.update({"fil_ph", "fil-ph", "tl", "tl_ph", "tl-ph", "tgl", "tgl_ph", "tgl-ph"}) # Expand across ISO variants / dataset-specific aliases. base_aliases = LANGUAGE_CODE_ALIASES.get(base, []) for alias in base_aliases: alias_norm = _norm_cfg_key(alias) out.update( { alias_norm, alias_norm.replace("_", "-"), _base_lang(alias_norm), } ) return list(out) def _pick_config_for_language(lang_code: str, available_configs: List[str]) -> Optional[str]: if not available_configs: return None infos = [] for cfg in available_configs: norm = _norm_cfg_key(cfg) infos.append((cfg, norm, _base_lang(norm))) candidates = set(_language_candidate_forms(lang_code)) code_norm = _norm_cfg_key(lang_code) code_base = _base_lang(code_norm) # 1) exact-ish match against candidate forms exact_hits = [t for t in infos if t[1] in candidates] if exact_hits: # Prefer exact code, then prefix match (e.g. en_us for en), then lexical. exact_hits.sort( key=lambda t: ( 0 if t[1] == code_norm else 1, 0 if t[1].startswith(code_base + "_") else 1, t[1], ) ) return exact_hits[0][0] # 2) base-language fallback base_hits = [t for t in infos if t[2] == code_base] if base_hits: # For Chinese variants, bias toward region when specified. def _rank(tup: Tuple[str, str, str]) -> Tuple[int, int, str]: norm = tup[1] if code_norm in {"zh_tw", "zhtw"}: zh_tw_bias = 0 if ("_tw" in norm or "hant" in norm) else 1 return (zh_tw_bias, 0 if norm.startswith(code_base + "_") else 1, norm) if code_norm == "zh": zh_cn_bias = 0 if ("_cn" in norm or "hans" in norm) else 1 return (zh_cn_bias, 0 if norm == "zh" else 1, norm) return (0, 0 if norm.startswith(code_base + "_") else 1, norm) base_hits.sort(key=_rank) return base_hits[0][0] return None def resolve_language_configs( language_codes: List[str], available_configs: List[str], ) -> Tuple[List[str], List[str], List[Tuple[str, str]]]: resolved: List[str] = [] missing: List[str] = [] mapping: List[Tuple[str, str]] = [] used = set() for code in language_codes: cfg = _pick_config_for_language(code, available_configs) if cfg is None: missing.append(code) continue mapping.append((code, cfg)) if cfg in used: continue resolved.append(cfg) used.add(cfg) return resolved, missing, mapping def resolve_configs(dataset_name: str, config_arg: str, all_configs: bool, max_configs: int) -> List[str]: available_configs = get_dataset_config_names(dataset_name, trust_remote_code=True) available_set = set(available_configs) preset = get_dataset_preset(dataset_name) alias_raw = preset.get("config_aliases", {}) alias_map = {k.lower(): v for k, v in alias_raw.items()} if all_configs: configs = available_configs else: raw_configs = [x.strip() for x in config_arg.split(",") if x.strip()] default_strategy = str(preset.get("default_strategy", "all_available")) # Fixed language datasets: map language codes -> actual config names. use_language_mapping = default_strategy == "fixed_languages" if use_language_mapping: if not raw_configs: raw_configs = list(preset.get("default_languages", [])) mapped_langs = [_map_config_alias(x, alias_map) for x in raw_configs] configs, missing_langs, mapping = resolve_language_configs(mapped_langs, available_configs) if mapping: mapping_str = ", ".join([f"{src}->{dst}" for src, dst in mapping]) print(f"[info] {dataset_name}: language->config mapping: {mapping_str}") if missing_langs: print( f"[warn] {dataset_name}: skip missing language configs: {missing_langs}" ) else: if not raw_configs: if default_strategy == "all_available": raw_configs = available_configs.copy() else: raw_configs = list(preset.get("default_configs", [])) mapped = [_map_config_alias(x, alias_map) for x in raw_configs] seen = set() configs = [] for cfg in mapped: if cfg and cfg not in seen: configs.append(cfg) seen.add(cfg) missing = [cfg for cfg in configs if cfg not in available_set] if missing: raise ValueError( "Unknown configs after alias mapping: " f"{missing}. Available configs include: {sorted(list(available_set))[:20]} ..." ) if not configs: raise ValueError("No configs resolved. Provide --configs or use --all-configs.") if max_configs > 0: configs = configs[:max_configs] return configs def dataset_column_names(raw_ds: any) -> List[str]: column_names = getattr(raw_ds, "column_names", None) if column_names: return list(column_names) features = getattr(raw_ds, "features", None) if features: try: return list(features.keys()) except Exception: pass return [] def resolve_text_column(dataset_name: str, column_names: List[str]) -> Optional[str]: preset = get_dataset_preset(dataset_name) preferred = list(preset.get("preferred_text_columns", [])) candidates = preferred + FALLBACK_TEXT_COLUMNS seen = set() deduped = [] for c in candidates: if c not in seen: deduped.append(c) seen.add(c) for c in deduped: if c in column_names: return c return None def resolve_im_start_id(tokenizer: any) -> int: im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>") if im_start_id is None or im_start_id == tokenizer.unk_token_id: if tokenizer.bos_token_id is not None: return int(tokenizer.bos_token_id) if tokenizer.eos_token_id is not None: return int(tokenizer.eos_token_id) return 0 return int(im_start_id) def ensure_front_prompt_token( asr_model: any, tokenizer: any, token_name: str = "<|latent|>", ) -> int: token_id = tokenizer.convert_tokens_to_ids(token_name) if token_id is None or token_id == tokenizer.unk_token_id: added = tokenizer.add_special_tokens({"additional_special_tokens": [token_name]}) if added > 0: asr_model.thinker.resize_token_embeddings(len(tokenizer)) token_id = tokenizer.convert_tokens_to_ids(token_name) if token_id is None or token_id < 0: raise RuntimeError(f"Failed to resolve front prompt token id for {token_name!r}.") return int(token_id) def extract_soft_prompt_tensor(ckpt: Dict[str, any]) -> Optional[torch.Tensor]: tensor = ckpt.get("soft_prompt_embed") if torch.is_tensor(tensor): return tensor state = ckpt.get("model_state_dict") if isinstance(state, dict): tensor = state.get("soft_prompt_embed") if torch.is_tensor(tensor): return tensor return None def load_checkpoint(path: str) -> Dict[str, any]: ckpt = torch.load(path, map_location="cpu") if isinstance(ckpt, dict): return ckpt return {"model_state_dict": ckpt} def is_path_file_or_dir(path: str) -> bool: return os.path.isfile(path) or os.path.isdir(path) def build_plain_bundle( model_id: str, device: str, dtype: torch.dtype, ) -> Tuple[LatentQwenASR, any]: asr_wrapper = Qwen3ASRModel.from_pretrained( model_id, dtype=dtype, device_map=device if device == "cuda" else None, ) asr_model = asr_wrapper.model processor = asr_wrapper.processor start_id = resolve_im_start_id(processor.tokenizer) model = LatentQwenASR( asr_model=asr_model, processor=processor, n_latent=0, nt_token_id=-1, lang_token_id=start_id, transcribe_token_id=start_id, freeze_base=False, use_latent=False, ).to(device) return model, processor def maybe_attach_peft_adapter( model: LatentQwenASR, adapter_path: str, mode_name: str, expected_lora_rank: Optional[int] = None, ) -> None: try: from peft import PeftConfig, PeftModel except Exception as e: raise RuntimeError( f"[{mode_name}] Loading PEFT adapter requires `peft` package: {e}" ) from e peft_cfg = PeftConfig.from_pretrained(adapter_path) peft_type = str(getattr(peft_cfg, "peft_type", "unknown")) print(f"[{mode_name}] adapter detected: peft_type={peft_type} path={adapter_path}") if expected_lora_rank is not None: rank = getattr(peft_cfg, "r", None) if rank is None: rank_pattern = getattr(peft_cfg, "rank_pattern", None) if isinstance(rank_pattern, dict) and rank_pattern: unique_ranks = sorted({int(v) for v in rank_pattern.values()}) if len(unique_ranks) == 1: rank = unique_ranks[0] if rank is None: print( f"[warn][{mode_name}] unable to verify LoRA rank from adapter config." ) elif int(rank) != int(expected_lora_rank): print( f"[warn][{mode_name}] expected LoRA rank={expected_lora_rank}, got rank={rank}." ) adapted = PeftModel.from_pretrained( model.thinker, adapter_path, is_trainable=False, ) model.base_model.thinker = adapted model.thinker = model.base_model.thinker def sanitize_generation_config(model: LatentQwenASR) -> None: gen_cfg = getattr(model.thinker, "generation_config", None) if gen_cfg is None: return # Keep decoding deterministic and avoid warnings from stale sampling flags. if hasattr(gen_cfg, "do_sample"): gen_cfg.do_sample = False if hasattr(gen_cfg, "temperature"): gen_cfg.temperature = 1.0 if hasattr(gen_cfg, "top_p"): gen_cfg.top_p = 1.0 if hasattr(gen_cfg, "typical_p"): gen_cfg.typical_p = 1.0 def build_base_model_bundle( model_id: str, device: str, dtype: torch.dtype, ) -> ModelBundle: print(f"\n[Load base_model] model={model_id}") ckpt: Dict[str, any] = {} model, processor = build_plain_bundle( model_id=model_id, device=device, dtype=dtype, ) sanitize_generation_config(model) model.eval() return ModelBundle(model=model, processor=processor, checkpoint=ckpt) def build_baseline_ft_bundle( model_id: str, checkpoint_path: str, device: str, dtype: torch.dtype, ) -> ModelBundle: print(f"\n[Load baseline_ft] model={model_id} ckpt={checkpoint_path}") ckpt = load_checkpoint(checkpoint_path) model, processor = build_plain_bundle( model_id=model_id, device=device, dtype=dtype, ) state = ckpt.get("model_state_dict", ckpt) missing, unexpected = model.load_state_dict(state, strict=False) if missing: print(f"[baseline_ft] missing keys: {len(missing)}") if unexpected: print(f"[baseline_ft] unexpected keys: {len(unexpected)}") sanitize_generation_config(model) model.eval() return ModelBundle(model=model, processor=processor, checkpoint=ckpt) def build_prompt_tuning_bundle( model_id: str, checkpoint_path: str, device: str, dtype: torch.dtype, ) -> ModelBundle: print(f"\n[Load prompt_tuning] model={model_id} ckpt={checkpoint_path}") if os.path.isdir(checkpoint_path): # Backward compatibility: old prompt-tuning PEFT adapters. model, processor = build_plain_bundle( model_id=model_id, device=device, dtype=dtype, ) maybe_attach_peft_adapter( model=model, adapter_path=checkpoint_path, mode_name="prompt_tuning", ) ckpt = { "adapter_path": checkpoint_path, "checkpoint_type": "peft_adapter", } else: ckpt = load_checkpoint(checkpoint_path) soft_prompt = extract_soft_prompt_tensor(ckpt) if soft_prompt is not None: asr_wrapper = Qwen3ASRModel.from_pretrained( model_id, dtype=dtype, device_map=device if device == "cuda" else None, ) asr_model = asr_wrapper.model processor = asr_wrapper.processor nt_id = ensure_front_prompt_token(asr_model, processor.tokenizer) start_id = resolve_im_start_id(processor.tokenizer) n_latent = int(ckpt.get("n_latent", int(soft_prompt.size(0)))) if n_latent <= 0: n_latent = int(soft_prompt.size(0)) if n_latent <= 0: raise ValueError("prompt_tuning checkpoint has empty soft prompt tensor.") model = LatentQwenASR( asr_model=asr_model, processor=processor, n_latent=n_latent, nt_token_id=nt_id, lang_token_id=start_id, transcribe_token_id=start_id, freeze_base=False, use_latent=False, use_soft_prompt=True, soft_prompt_init_mode="random", ).to(device) prompt = soft_prompt.detach() if prompt.dim() != 2: raise ValueError( f"soft_prompt_embed must be 2D [n_latent, hidden], got shape={tuple(prompt.shape)}" ) if prompt.size(0) < n_latent: reps = (n_latent + prompt.size(0) - 1) // prompt.size(0) prompt = prompt.repeat((reps, 1)) prompt = prompt[:n_latent] model.soft_prompt_embed.data.copy_( prompt.to( device=model.soft_prompt_embed.device, dtype=model.soft_prompt_embed.dtype, ) ) print(f"[prompt_tuning] loaded front soft prompt: n_latent={n_latent}") else: # Fallback for legacy `.pth` checkpoints that saved full wrapper state. model, processor = build_plain_bundle( model_id=model_id, device=device, dtype=dtype, ) state = ckpt.get("model_state_dict", ckpt) missing, unexpected = model.load_state_dict(state, strict=False) if missing: print(f"[prompt_tuning] missing keys: {len(missing)}") if unexpected: print(f"[prompt_tuning] unexpected keys: {len(unexpected)}") sanitize_generation_config(model) model.eval() return ModelBundle(model=model, processor=processor, checkpoint=ckpt) def build_lora_r16_bundle( model_id: str, checkpoint_path: str, device: str, dtype: torch.dtype, ) -> ModelBundle: print(f"\n[Load lora_r16] model={model_id} ckpt={checkpoint_path}") model, processor = build_plain_bundle( model_id=model_id, device=device, dtype=dtype, ) if os.path.isdir(checkpoint_path): maybe_attach_peft_adapter( model=model, adapter_path=checkpoint_path, mode_name="lora_r16", expected_lora_rank=None, ) ckpt = { "adapter_path": checkpoint_path, "checkpoint_type": "peft_adapter", } else: ckpt = load_checkpoint(checkpoint_path) state = ckpt.get("model_state_dict", ckpt) missing, unexpected = model.load_state_dict(state, strict=False) if missing: print(f"[lora_r16] missing keys: {len(missing)}") if unexpected: print(f"[lora_r16] unexpected keys: {len(unexpected)}") sanitize_generation_config(model) model.eval() return ModelBundle(model=model, processor=processor, checkpoint=ckpt) def build_latent_bundle( model_id: str, checkpoint_path: str, n_latent_override: int, device: str, dtype: torch.dtype, ) -> ModelBundle: print(f"\n[Load latent_reasoning] model={model_id} ckpt={checkpoint_path}") ckpt = load_checkpoint(checkpoint_path) n_latent = int(ckpt.get("n_latent", 4)) if n_latent_override >= 0: n_latent = int(n_latent_override) if n_latent <= 0: raise ValueError(f"Invalid latent token count: {n_latent}") asr_wrapper = Qwen3ASRModel.from_pretrained( model_id, dtype=dtype, device_map=device if device == "cuda" else None, ) asr_model = asr_wrapper.model processor = asr_wrapper.processor token_id = ensure_front_prompt_token(asr_model, processor.tokenizer) start_id = resolve_im_start_id(processor.tokenizer) model = LatentQwenASR( asr_model=asr_model, processor=processor, n_latent=n_latent, nt_token_id=int(token_id), lang_token_id=start_id, transcribe_token_id=start_id, freeze_base=False, use_latent=True, latent_use_bounded_delta=bool(ckpt.get("latent_use_bounded_delta", True)), latent_use_injection_gate=bool(ckpt.get("latent_use_injection_gate", True)), latent_use_embedding_anchor=bool(ckpt.get("latent_use_embedding_anchor", True)), ).to(device) # Prefer latent-adapter payload format from train.py. if "init_proj" in ckpt: model.init_proj.load_state_dict(ckpt["init_proj"], strict=True) model.delta_proj.load_state_dict(ckpt["delta_proj"], strict=True) model.step_proj.load_state_dict(ckpt["step_proj"], strict=True) if "step_embed" in ckpt: model.step_embed.data.copy_( ckpt["step_embed"].to( device=model.step_embed.device, dtype=model.step_embed.dtype, ) ) if "log_scale" in ckpt: model.log_scale.data.copy_( ckpt["log_scale"].to( device=model.log_scale.device, dtype=model.log_scale.dtype, ) ) if "value_head" in ckpt: model.value_head.load_state_dict(ckpt["value_head"], strict=True) if "injection_gate" in ckpt: model.injection_gate.load_state_dict(ckpt["injection_gate"], strict=True) else: # Fallback: checkpoint may contain full wrapper state_dict. state = ckpt.get("model_state_dict", ckpt) missing, unexpected = model.load_state_dict(state, strict=False) if missing: print(f"[latent_reasoning] missing keys: {len(missing)}") if unexpected: print(f"[latent_reasoning] unexpected keys: {len(unexpected)}") sanitize_generation_config(model) model.eval() return ModelBundle(model=model, processor=processor, checkpoint=ckpt) _global_normalizer = None _global_normalizer_kind = "english" def configure_text_normalizer(kind: str) -> None: global _global_normalizer, _global_normalizer_kind kind = (kind or "english").strip().lower() if kind not in {"english", "basic", "simple"}: raise ValueError(f"Unknown text normalizer: {kind}") if kind != _global_normalizer_kind: _global_normalizer = None _global_normalizer_kind = kind def simple_text_normalize(text: str) -> str: text = unicodedata.normalize("NFKC", text or "") text = text.lower() text = "".join( " " if unicodedata.category(ch).startswith("P") else ch for ch in text ) return re.sub(r"\s+", " ", text).strip() class _SimpleTextNormalizer: def __call__(self, text: str) -> str: return simple_text_normalize(text) def build_text_normalizer(kind: str): if kind == "simple": return _SimpleTextNormalizer() if kind == "basic": try: from transformers.models.whisper.english_normalizer import BasicTextNormalizer return BasicTextNormalizer() except Exception: return _SimpleTextNormalizer() try: # We use the standard Whisper English normalizer for English-heavy legacy runs. from transformers.models.whisper.english_normalizer import EnglishTextNormalizer import json, urllib.request, os mapping_path = os.path.join(os.path.dirname(__file__), "english_normalizer.json") if not os.path.exists(mapping_path): url = "https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/models/whisper/english_normalizer.json" req = urllib.request.Request(url) with urllib.request.urlopen(req) as response: data = json.loads(response.read().decode("utf-8")) with open(mapping_path, "w") as f: json.dump(data, f) else: with open(mapping_path, "r") as f: data = json.load(f) return EnglishTextNormalizer(data) except Exception: try: from transformers.models.whisper.english_normalizer import BasicTextNormalizer return BasicTextNormalizer() except Exception: return _SimpleTextNormalizer() def normalize_text(text: str) -> str: global _global_normalizer if _global_normalizer is None: _global_normalizer = build_text_normalizer(_global_normalizer_kind) normalized = _global_normalizer(text) if isinstance(normalized, str): if _global_normalizer_kind == "english": # LibriSpeech specific normalizations: expand common contractions # that Whisper leaves as "m" or "re". contractions = { r"\bi m\b": "i am", r"\byou re\b": "you are", r"\bwe re\b": "we are", r"\bthey re\b": "they are", r"\bhe s\b": "he is", r"\bshe s\b": "she is", r"\bit s\b": "it is", r"\bthat s\b": "that is", r"\bwho s\b": "who is", r"\bwhat s\b": "what is", r"\bthere s\b": "there is", r"\bhere s\b": "here is", r"\bwhere s\b": "where is", r"\blet s\b": "let us", r"\bhow s\b": "how is", # Standard possessive stripping (boy's -> boy s -> boys) r"\b s\b": "s", } for pat, repl in contractions.items(): normalized = re.sub(pat, repl, normalized) return re.sub(r"\s+", " ", normalized).strip() return "" def clean_prediction(text: str) -> str: text = text or "" text = re.sub(r"language\s+\w+", "", text, flags=re.IGNORECASE) if "" in text: text = text.split("", 1)[-1] return text.strip() @torch.no_grad() def add_noise_at_snr(waveform: np.ndarray, snr_db: float) -> np.ndarray: """Add Gaussian noise to a waveform at a target Signal-to-Noise Ratio (dB). SNR_dB = 10 * log10(P_signal / P_noise) => P_noise = P_signal / 10^(SNR_dB / 10) => std_noise = sqrt(P_noise) """ signal_power = np.mean(waveform ** 2) if signal_power < 1e-10: # silence guard return waveform noise_power = signal_power / (10.0 ** (snr_db / 10.0)) noise = np.random.normal(0, np.sqrt(noise_power), waveform.shape) return (waveform + noise).astype(waveform.dtype) def evaluate_config( model: LatentQwenASR, processor: any, dataset_name: str, config_name: str, split: str, use_baseline: bool, mode_name: str, max_samples: int, max_new_tokens: int, print_samples: int, language_hint: Optional[str], prompt_template: str, disable_language_hint: bool, num_beams: int = 1, dynamic_halt_threshold: float = 0.0, snr_db: Optional[float] = None, print_gate_skips: bool = False, streaming: bool = False, ) -> Dict[str, any]: raw_ds = load_dataset( dataset_name, config_name, split=split, trust_remote_code=True, streaming=streaming, ) known_total: Optional[int] try: known_total = len(raw_ds) except TypeError: known_total = None total = known_total if known_total is not None else 0 if max_samples > 0: total = min(total, max_samples) if known_total is not None else max_samples column_names = dataset_column_names(raw_ds) text_column = resolve_text_column(dataset_name, column_names) audio_column = "audio" if "audio" in column_names else None if audio_column is None and not streaming: raise ValueError( f"Dataset {dataset_name}/{config_name} does not have an 'audio' column. " f"Columns: {column_names}" ) target_sampling_rate = int(getattr(processor.feature_extractor, "sampling_rate", 0) or 0) if target_sampling_rate > 0 and audio_column is not None: raw_ds = raw_ds.cast_column(audio_column, Audio(sampling_rate=target_sampling_rate)) preds: List[str] = [] refs: List[str] = [] shown = 0 skipped = 0 reasoned = 0 # Value Head aggregate statistics collectors all_v_preds: List[List[float]] = [] # per-sample list of v_pred values all_deq_iters: List[int] = [] # per-sample thinking steps taken all_gates: List[float] = [] # per-sample gate values all_scales: List[List[float]] = [] # per-sample scale values n_skipped_by_gate = 0 # count of N=0 (gate skip) effective_language_hint = None if disable_language_hint else language_hint if effective_language_hint: try: prompt_text = prompt_template.format(language=effective_language_hint) except Exception: prompt_text = f"Transcribe the {effective_language_hint} audio into text." else: prompt_text = "Transcribe the audio into text." sample_iter = iter(raw_ds) if max_samples > 0: sample_iter = itertools.islice(sample_iter, max_samples) pbar = tqdm(sample_iter, total=total if total > 0 else None, desc=f"{config_name}:{mode_name}", leave=False) seen = 0 for i, sample in enumerate(pbar): seen += 1 if (text_column is None or audio_column is None) and isinstance(sample, dict): sample_columns = list(sample.keys()) if text_column is None: text_column = resolve_text_column(dataset_name, sample_columns) if audio_column is None and "audio" in sample_columns: audio_column = "audio" if audio_column is None: skipped += 1 continue if text_column is not None: ref = sample.get(text_column, "") else: ref = "" for cand in FALLBACK_TEXT_COLUMNS: v = sample.get(cand, "") if isinstance(v, str) and v.strip(): ref = v break if not isinstance(ref, str) or not ref.strip(): skipped += 1 continue audio = sample.get(audio_column) if not isinstance(audio, dict) or "array" not in audio or "sampling_rate" not in audio: skipped += 1 continue audio_array = np.array(audio["array"], dtype=np.float64) if snr_db is not None: audio_array = add_noise_at_snr(audio_array, snr_db) 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, ) feats = torch.tensor(feat_out.input_features[0], dtype=target_dtype, device=model.base_model.device).unsqueeze(0) n_frames = feats.size(-1) if getattr(feat_out, "attention_mask", None) is not None: # Coerce mask to match the Mel frame length (same logic as data.py) 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) # Pad or truncate to n_frames 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=model.base_model.device).unsqueeze(0) else: feature_attention_mask = torch.ones((1, n_frames), dtype=torch.long, device=model.base_model.device) # Guard: if effective feature length is too short for audio tower, skip sample effective_len = feature_attention_mask.sum().item() if effective_len < 10: skipped += 1 continue gen_kwargs = { "feature_attention_mask": 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": num_beams, "language_hint": effective_language_hint, "prompt_text": prompt_text, "dynamic_halt_threshold": dynamic_halt_threshold, } gen_output = model.generate(feats, **gen_kwargs) if isinstance(gen_output, tuple): gen_ids = gen_output[0] stats = gen_output[1] if isinstance(gen_output[1], dict) else (gen_output[2] if len(gen_output)>2 and isinstance(gen_output[2], dict) else {}) else: gen_ids = gen_output 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] pred_raw = processor.tokenizer.decode(ids, skip_special_tokens=True) pred = clean_prediction(pred_raw) pred_norm = normalize_text(pred) ref_norm = normalize_text(ref) if not ref_norm: skipped += 1 continue preds.append(pred_norm) refs.append(ref_norm) # Always print if it was skipped (N=0 from Gate) OR if under the print limit is_skipped_lr = stats.get("skipped", False) if mode_name == "latent_reasoning": if not is_skipped_lr: reasoned += 1 if shown < print_samples or (print_gate_skips and is_skipped_lr): if not is_skipped_lr: shown += 1 print(f"\n[{config_name} | {mode_name} | sample {shown if not is_skipped_lr else 'SKIPPED'}]") print(f" Ref : {ref.strip()}") print(f" Pred: {pred}") if "v_preds" in stats and stats["v_preds"] is not None: v_vals = stats["v_preds"][0].tolist() v_str = " | ".join([f"{v:.4f}" for v in v_vals]) iters_t = stats.get("deq_iters", None) deq_iters = int(iters_t.item()) if iters_t is not None and hasattr(iters_t, "item") else -1 if is_skipped_lr: print(f" v_preds: [{v_str}] -> SKIPPED (N=0)") else: print(f" v_preds: [{v_str}] -> {deq_iters} step(s)") # --- Collect Value Head stats for aggregate summary --- if mode_name == "latent_reasoning": iters_t = stats.get("deq_iters", None) deq_n = int(iters_t.item()) if iters_t is not None and hasattr(iters_t, "item") else 0 all_deq_iters.append(deq_n) if is_skipped_lr: n_skipped_by_gate += 1 if "v_preds" in stats and stats["v_preds"] is not None: all_v_preds.append(stats["v_preds"][0].tolist()) if "gate" in stats and stats["gate"] is not None: g = stats["gate"] all_gates.append(float(g[0].item()) if hasattr(g, "item") or g.dim() > 0 else float(g)) if "scales" in stats and stats["scales"] is not None: s = stats["scales"] all_scales.append(s.tolist() if hasattr(s, "tolist") else list(s)) # --- Print Value Head aggregate summary --- if mode_name == "latent_reasoning" and all_deq_iters: import collections n_total_eval = len(all_deq_iters) n_dist = collections.Counter(all_deq_iters) print(f"\n{'='*60}") print(f"[Value-Head Summary] {config_name} | {mode_name}") print(f"{'='*60}") print(f" Total samples evaluated : {n_total_eval}") print(f" Skipped by gate (N=0) : {n_skipped_by_gate} ({100*n_skipped_by_gate/max(1,n_total_eval):.1f}%)") print(f" Reasoned (N>0) : {n_total_eval - n_skipped_by_gate} ({100*(n_total_eval - n_skipped_by_gate)/max(1,n_total_eval):.1f}%)") print(f" --- N-step distribution ---") for k in sorted(n_dist.keys()): pct = 100 * n_dist[k] / n_total_eval bar = '█' * int(pct / 2) print(f" N={k}: {n_dist[k]:5d} ({pct:5.1f}%) {bar}") if all_v_preds: flat_v = [v for vs in all_v_preds for v in vs] if flat_v: import numpy as _np v_arr = _np.array(flat_v) print(f" --- v_pred statistics (all steps, all samples) ---") print(f" count : {len(v_arr)}") print(f" mean : {v_arr.mean():.4f}") print(f" std : {v_arr.std():.4f}") print(f" min : {v_arr.min():.4f}") print(f" max : {v_arr.max():.4f}") print(f" median: {float(_np.median(v_arr)):.4f}") # Histogram buckets edges = [-1.0, -0.8, -0.5, -0.2, 0.0, 0.2, 0.5, 0.8, 1.0] print(f" --- v_pred histogram ---") for j in range(len(edges) - 1): lo, hi = edges[j], edges[j + 1] cnt = int(((v_arr >= lo) & (v_arr < hi)).sum()) pct = 100 * cnt / len(v_arr) bar = '▓' * int(pct / 2) print(f" [{lo:+.1f}, {hi:+.1f}): {cnt:5d} ({pct:5.1f}%) {bar}") # Also count == 1.0 cnt_one = int((v_arr >= 0.8).sum()) # First-step v_pred (initial gate signal) first_v = [vs[0] for vs in all_v_preds if vs] if first_v: fv = _np.array(first_v) print(f" --- v_pred[0] (initial gate signal) ---") print(f" mean : {fv.mean():.4f}") print(f" std : {fv.std():.4f}") print(f" < 0 : {int((fv < 0).sum()):5d} ({100*(fv < 0).mean():.1f}%)") print(f" >= 0 : {int((fv >= 0).sum()):5d} ({100*(fv >= 0).mean():.1f}%)") if all_gates: import numpy as _np g_arr = _np.array(all_gates) print(f" --- Gate confidence (mapped to [0,1]) ---") print(f" mean : {g_arr.mean():.4f}") print(f" std : {g_arr.std():.4f}") print(f" min : {g_arr.min():.4f}") print(f" max : {g_arr.max():.4f}") if all_scales: import numpy as _np s_arr = _np.array(all_scales) if s_arr.ndim == 2: print(f" --- Delta scales per step ---") for si in range(s_arr.shape[1]): col = s_arr[:, si] print(f" step {si}: mean={col.mean():.4f} std={col.std():.4f}") elif s_arr.ndim == 1: print(f" --- Delta scales ---") print(f" mean={s_arr.mean():.4f} std={s_arr.std():.4f}") print(f"{'='*60}\n") if not refs: return { "config": config_name, "mode": mode_name, "language_hint": effective_language_hint, "text_column": text_column, "samples_total": total if total > 0 else seen, "samples_used": 0, "samples_skipped": skipped, "samples_reasoned": reasoned, "wer": 1.0, "cer": 1.0, } return { "config": config_name, "mode": mode_name, "language_hint": effective_language_hint, "text_column": text_column, "samples_total": total if total > 0 else seen, "samples_used": len(refs), "samples_skipped": skipped, "samples_reasoned": reasoned, "wer": float(wer(refs, preds)), "cer": float(cer(refs, preds)), } def weighted_avg(rows: List[Dict[str, any]], metric_key: str) -> Optional[float]: numer = 0.0 denom = 0 for row in rows: n = int(row.get("samples_used", 0)) if n <= 0: continue m = row.get(metric_key) if m is None: continue numer += float(m) * n denom += n if denom == 0: return None return numer / float(denom) def print_summary_table(rows: List[Dict[str, any]]) -> None: header = ( f"{'Config':<12} {'N':>7} " f"{'Raw-WER':>10} {'FT-WER':>10} {'PT-WER':>10} {'LoRA-WER':>10} {'Lat-WER':>10} " f"{'Raw-CER':>10} {'FT-CER':>10} {'PT-CER':>10} {'LoRA-CER':>10} {'Lat-CER':>10}" ) print("\n" + header) print("-" * len(header)) for row in rows: cfg = row["config"] n = row.get("samples_used", 0) raw_wer = row.get("base_model_wer") ft_wer = row.get("baseline_ft_wer") pt_wer = row.get("prompt_tuning_wer") lora_wer = row.get("lora_r16_wer") lat_wer = row.get("latent_reasoning_wer") raw_cer = row.get("base_model_cer") ft_cer = row.get("baseline_ft_cer") pt_cer = row.get("prompt_tuning_cer") lora_cer = row.get("lora_r16_cer") lat_cer = row.get("latent_reasoning_cer") print( f"{cfg:<12} {n:>7d} " f"{(f'{raw_wer:.4f}' if raw_wer is not None else '-'):>10} " f"{(f'{ft_wer:.4f}' if ft_wer is not None else '-'):>10} " f"{(f'{pt_wer:.4f}' if pt_wer is not None else '-'):>10} " f"{(f'{lora_wer:.4f}' if lora_wer is not None else '-'):>10} " f"{(f'{lat_wer:.4f}' if lat_wer is not None else '-'):>10} " f"{(f'{raw_cer:.4f}' if raw_cer is not None else '-'):>10} " f"{(f'{ft_cer:.4f}' if ft_cer is not None else '-'):>10} " f"{(f'{pt_cer:.4f}' if pt_cer is not None else '-'):>10} " f"{(f'{lora_cer:.4f}' if lora_cer is not None else '-'):>10} " f"{(f'{lat_cer:.4f}' if lat_cer is not None else '-'):>10}" ) def release_model(bundle: Optional[ModelBundle]) -> None: if bundle is None: return del bundle.model gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def main() -> None: args = parse_args() set_seed(args.seed) configure_text_normalizer(args.text_normalizer) args.baseline_ckpt = (args.baseline_ckpt or "").strip() args.prompt_tuning_ckpt = (args.prompt_tuning_ckpt or "").strip() args.lora_r16_ckpt = (args.lora_r16_ckpt or "").strip() args.latent_ckpt = (args.latent_ckpt or "").strip() run_base_model = not bool(args.skip_base_model) run_baseline_ft = bool(args.baseline_ckpt) and not bool(args.skip_baseline_ft) run_prompt_tuning = bool(args.prompt_tuning_ckpt) and not bool(args.skip_prompt_tuning) run_lora_r16 = bool(args.lora_r16_ckpt) and not bool(args.skip_lora_r16) run_latent_reasoning = bool(args.latent_ckpt) if not ( run_base_model or run_baseline_ft or run_prompt_tuning or run_lora_r16 or run_latent_reasoning ): raise ValueError( "No evaluation mode enabled. Provide at least one of --baseline-ckpt, " "--prompt-tuning-ckpt, --lora-r16-ckpt, --latent-ckpt, " "or remove --skip-base-model." ) if run_baseline_ft and not os.path.isfile(args.baseline_ckpt): raise FileNotFoundError( f"--baseline-ckpt expects a .pth file, got: {args.baseline_ckpt}" ) if run_latent_reasoning and not os.path.isfile(args.latent_ckpt): raise FileNotFoundError( f"--latent-ckpt expects a .pth file, got: {args.latent_ckpt}" ) if run_prompt_tuning and not is_path_file_or_dir(args.prompt_tuning_ckpt): raise FileNotFoundError( f"--prompt-tuning-ckpt path not found: {args.prompt_tuning_ckpt}" ) if run_lora_r16 and not is_path_file_or_dir(args.lora_r16_ckpt): raise FileNotFoundError( f"--lora-r16-ckpt path not found: {args.lora_r16_ckpt}" ) dataset_name = resolve_dataset_name(args.dataset_name) device = choose_device(args.device) dtype = choose_dtype(args.dtype, device) configs = resolve_configs(dataset_name, args.configs, args.all_configs, args.max_configs) print("=== Evaluation Setup ===") print(f"Device: {device}") print(f"Dtype : {dtype}") print(f"Model : {args.model_id}") print(f"Data : {dataset_name} ({args.split})") print(f"Configs ({len(configs)}): {configs}") print(f"Max samples/config: {'all' if args.max_samples_per_config <= 0 else args.max_samples_per_config}") print(f"Streaming: {'yes' if args.streaming else 'no'}") print(f"Text normalizer: {args.text_normalizer}") print(f"Run base_model: {'yes' if run_base_model else 'no'}") print(f"Run baseline_ft: {'yes' if run_baseline_ft else 'no'}") print(f"Run prompt_tuning: {'yes' if run_prompt_tuning else 'no'}") print(f"Run lora_r16: {'yes' if run_lora_r16 else 'no'}") print(f"Run latent_reasoning: {'yes' if run_latent_reasoning else 'no'}") print(f"Language hint: {'off' if args.no_language_hint else 'on'}") if not args.no_language_hint: print(f"Prompt template: {args.prompt_template}") print(f"Num beams: {args.num_beams}") print(f"Print gate skips: {'yes' if args.print_gate_skips else 'no'}") base_model_rows: Dict[str, Dict[str, any]] = {} baseline_ft_rows: Dict[str, Dict[str, any]] = {} prompt_tuning_rows: Dict[str, Dict[str, any]] = {} lora_r16_rows: Dict[str, Dict[str, any]] = {} latent_reasoning_rows: Dict[str, Dict[str, any]] = {} base_model_bundle: Optional[ModelBundle] = None baseline_ft_bundle: Optional[ModelBundle] = None prompt_tuning_bundle: Optional[ModelBundle] = None lora_r16_bundle: Optional[ModelBundle] = None latent_reasoning_bundle: Optional[ModelBundle] = None if run_base_model: try: base_model_bundle = build_base_model_bundle( model_id=args.model_id, device=device, dtype=dtype, ) for cfg in configs: language_hint = resolve_language_hint(dataset_name, cfg) out = evaluate_config( model=base_model_bundle.model, processor=base_model_bundle.processor, dataset_name=dataset_name, config_name=cfg, split=args.split, use_baseline=True, mode_name="base_model", max_samples=args.max_samples_per_config, max_new_tokens=args.max_new_tokens, print_samples=args.print_samples, language_hint=language_hint, prompt_template=args.prompt_template, disable_language_hint=args.no_language_hint, num_beams=args.num_beams, dynamic_halt_threshold=args.dynamic_halt_threshold, snr_db=args.snr_db[0] if args.snr_db else None, print_gate_skips=args.print_gate_skips, streaming=args.streaming, ) base_model_rows[cfg] = out print( f"[BaseModel][{cfg}] " f"lang={out.get('language_hint') or 'auto'} " f"text_col={out.get('text_column') or '-'} " f"WER={out['wer']:.4f} CER={out['cer']:.4f} " f"n={out['samples_used']} skip={out['samples_skipped']}" ) finally: release_model(base_model_bundle) else: print("\n[Skip] base_model evaluation skipped (--skip-base-model).") if run_baseline_ft: try: baseline_ft_bundle = build_baseline_ft_bundle( model_id=args.model_id, checkpoint_path=args.baseline_ckpt, device=device, dtype=dtype, ) for cfg in configs: language_hint = resolve_language_hint(dataset_name, cfg) out = evaluate_config( model=baseline_ft_bundle.model, processor=baseline_ft_bundle.processor, dataset_name=dataset_name, config_name=cfg, split=args.split, use_baseline=True, mode_name="baseline_ft", max_samples=args.max_samples_per_config, max_new_tokens=args.max_new_tokens, print_samples=args.print_samples, language_hint=language_hint, prompt_template=args.prompt_template, disable_language_hint=args.no_language_hint, num_beams=args.num_beams, dynamic_halt_threshold=args.dynamic_halt_threshold, snr_db=args.snr_db[0] if args.snr_db else None, print_gate_skips=args.print_gate_skips, streaming=args.streaming, ) baseline_ft_rows[cfg] = out print( f"[BaselineFT][{cfg}] " f"lang={out.get('language_hint') or 'auto'} " f"text_col={out.get('text_column') or '-'} " f"WER={out['wer']:.4f} CER={out['cer']:.4f} " f"n={out['samples_used']} skip={out['samples_skipped']}" ) finally: release_model(baseline_ft_bundle) else: print("\n[Skip] baseline_ft evaluation skipped (no --baseline-ckpt).") if run_prompt_tuning: try: prompt_tuning_bundle = build_prompt_tuning_bundle( model_id=args.model_id, checkpoint_path=args.prompt_tuning_ckpt, device=device, dtype=dtype, ) for cfg in configs: language_hint = resolve_language_hint(dataset_name, cfg) out = evaluate_config( model=prompt_tuning_bundle.model, processor=prompt_tuning_bundle.processor, dataset_name=dataset_name, config_name=cfg, split=args.split, use_baseline=False, mode_name="prompt_tuning", max_samples=args.max_samples_per_config, max_new_tokens=args.max_new_tokens, print_samples=args.print_samples, language_hint=language_hint, prompt_template=args.prompt_template, disable_language_hint=args.no_language_hint, num_beams=args.num_beams, dynamic_halt_threshold=args.dynamic_halt_threshold, snr_db=args.snr_db[0] if args.snr_db else None, print_gate_skips=args.print_gate_skips, streaming=args.streaming, ) prompt_tuning_rows[cfg] = out print( f"[PromptTuning][{cfg}] " f"text_col={out.get('text_column') or '-'} " f"WER={out['wer']:.4f} CER={out['cer']:.4f} " f"n={out['samples_used']} skip={out['samples_skipped']}" ) finally: release_model(prompt_tuning_bundle) else: print("\n[Skip] prompt_tuning evaluation skipped (no --prompt-tuning-ckpt).") if run_lora_r16: try: lora_r16_bundle = build_lora_r16_bundle( model_id=args.model_id, checkpoint_path=args.lora_r16_ckpt, device=device, dtype=dtype, ) for cfg in configs: language_hint = resolve_language_hint(dataset_name, cfg) out = evaluate_config( model=lora_r16_bundle.model, processor=lora_r16_bundle.processor, dataset_name=dataset_name, config_name=cfg, split=args.split, use_baseline=True, mode_name="lora_r16", max_samples=args.max_samples_per_config, max_new_tokens=args.max_new_tokens, print_samples=args.print_samples, language_hint=language_hint, prompt_template=args.prompt_template, disable_language_hint=args.no_language_hint, num_beams=args.num_beams, dynamic_halt_threshold=args.dynamic_halt_threshold, snr_db=args.snr_db[0] if args.snr_db else None, print_gate_skips=args.print_gate_skips, streaming=args.streaming, ) lora_r16_rows[cfg] = out print( f"[LoRA-R16][{cfg}] " f"text_col={out.get('text_column') or '-'} " f"WER={out['wer']:.4f} CER={out['cer']:.4f} " f"n={out['samples_used']} skip={out['samples_skipped']}" ) finally: release_model(lora_r16_bundle) else: print("\n[Skip] lora_r16 evaluation skipped (no --lora-r16-ckpt).") if run_latent_reasoning: try: latent_reasoning_bundle = build_latent_bundle( model_id=args.model_id, checkpoint_path=args.latent_ckpt, n_latent_override=args.n_latent, device=device, dtype=dtype, ) for cfg in configs: language_hint = resolve_language_hint(dataset_name, cfg) out = evaluate_config( model=latent_reasoning_bundle.model, processor=latent_reasoning_bundle.processor, dataset_name=dataset_name, config_name=cfg, split=args.split, use_baseline=False, mode_name="latent_reasoning", max_samples=args.max_samples_per_config, max_new_tokens=args.max_new_tokens, print_samples=args.print_samples, language_hint=language_hint, prompt_template=args.prompt_template, disable_language_hint=args.no_language_hint, num_beams=args.num_beams, dynamic_halt_threshold=args.dynamic_halt_threshold, snr_db=args.snr_db[0] if args.snr_db else None, print_gate_skips=args.print_gate_skips, streaming=args.streaming, ) latent_reasoning_rows[cfg] = out n_used = out['samples_used'] n_reasoned = out.get('samples_reasoned', 0) reason_pct = (n_reasoned / n_used * 100) if n_used > 0 else 0.0 print( f"[LatentReasoning][{cfg}] " f"lang={out.get('language_hint') or 'auto'} " f"text_col={out.get('text_column') or '-'} " f"WER={out['wer']:.4f} CER={out['cer']:.4f} " f"n={n_used} skip={out['samples_skipped']} " f"reasoned={n_reasoned}/{n_used} ({reason_pct:.1f}%)" ) finally: release_model(latent_reasoning_bundle) else: print("\n[Skip] latent_reasoning evaluation skipped (no --latent-ckpt).") merged_rows: List[Dict[str, any]] = [] for cfg in configs: base = base_model_rows.get(cfg, {}) ft = baseline_ft_rows.get(cfg, {}) prompt = prompt_tuning_rows.get(cfg, {}) lora = lora_r16_rows.get(cfg, {}) lat = latent_reasoning_rows.get(cfg, {}) text_col = ( base.get("text_column") or ft.get("text_column") or prompt.get("text_column") or lora.get("text_column") or lat.get("text_column") ) row = { "config": cfg, "samples_total": int( max( base.get("samples_total", 0), ft.get("samples_total", 0), prompt.get("samples_total", 0), lora.get("samples_total", 0), lat.get("samples_total", 0), ) ), "samples_used": int( max( base.get("samples_used", 0), ft.get("samples_used", 0), prompt.get("samples_used", 0), lora.get("samples_used", 0), lat.get("samples_used", 0), ) ), "text_column": text_col, "base_model_wer": base.get("wer"), "base_model_cer": base.get("cer"), "baseline_ft_wer": ft.get("wer"), "baseline_ft_cer": ft.get("cer"), "prompt_tuning_wer": prompt.get("wer"), "prompt_tuning_cer": prompt.get("cer"), "lora_r16_wer": lora.get("wer"), "lora_r16_cer": lora.get("cer"), "latent_reasoning_wer": lat.get("wer"), "latent_reasoning_cer": lat.get("cer"), "baseline_ft_vs_base_model_wer": None if base.get("wer") is None or ft.get("wer") is None else float(base["wer"] - ft["wer"]), "baseline_ft_vs_base_model_cer": None if base.get("cer") is None or ft.get("cer") is None else float(base["cer"] - ft["cer"]), "prompt_tuning_vs_base_model_wer": None if base.get("wer") is None or prompt.get("wer") is None else float(base["wer"] - prompt["wer"]), "prompt_tuning_vs_base_model_cer": None if base.get("cer") is None or prompt.get("cer") is None else float(base["cer"] - prompt["cer"]), "lora_r16_vs_base_model_wer": None if base.get("wer") is None or lora.get("wer") is None else float(base["wer"] - lora["wer"]), "lora_r16_vs_base_model_cer": None if base.get("cer") is None or lora.get("cer") is None else float(base["cer"] - lora["cer"]), "latent_reasoning_vs_base_model_wer": None if base.get("wer") is None or lat.get("wer") is None else float(base["wer"] - lat["wer"]), "latent_reasoning_vs_base_model_cer": None if base.get("cer") is None or lat.get("cer") is None else float(base["cer"] - lat["cer"]), "latent_reasoning_vs_baseline_ft_wer": None if ft.get("wer") is None or lat.get("wer") is None else float(ft["wer"] - lat["wer"]), "latent_reasoning_vs_baseline_ft_cer": None if ft.get("cer") is None or lat.get("cer") is None else float(ft["cer"] - lat["cer"]), "latent_reasoning_vs_prompt_tuning_wer": None if prompt.get("wer") is None or lat.get("wer") is None else float(prompt["wer"] - lat["wer"]), "latent_reasoning_vs_prompt_tuning_cer": None if prompt.get("cer") is None or lat.get("cer") is None else float(prompt["cer"] - lat["cer"]), "latent_reasoning_vs_lora_r16_wer": None if lora.get("wer") is None or lat.get("wer") is None else float(lora["wer"] - lat["wer"]), "latent_reasoning_vs_lora_r16_cer": None if lora.get("cer") is None or lat.get("cer") is None else float(lora["cer"] - lat["cer"]), } merged_rows.append(row) print_summary_table(merged_rows) base_list = [base_model_rows[c] for c in configs if c in base_model_rows] ft_list = [baseline_ft_rows[c] for c in configs if c in baseline_ft_rows] prompt_list = [prompt_tuning_rows[c] for c in configs if c in prompt_tuning_rows] lora_list = [lora_r16_rows[c] for c in configs if c in lora_r16_rows] lat_list = [latent_reasoning_rows[c] for c in configs if c in latent_reasoning_rows] summary = { "base_model_weighted_wer": weighted_avg(base_list, "wer"), "base_model_weighted_cer": weighted_avg(base_list, "cer"), "baseline_ft_weighted_wer": weighted_avg(ft_list, "wer"), "baseline_ft_weighted_cer": weighted_avg(ft_list, "cer"), "prompt_tuning_weighted_wer": weighted_avg(prompt_list, "wer"), "prompt_tuning_weighted_cer": weighted_avg(prompt_list, "cer"), "lora_r16_weighted_wer": weighted_avg(lora_list, "wer"), "lora_r16_weighted_cer": weighted_avg(lora_list, "cer"), "latent_reasoning_weighted_wer": weighted_avg(lat_list, "wer"), "latent_reasoning_weighted_cer": weighted_avg(lat_list, "cer"), } if summary["base_model_weighted_wer"] is not None and summary["baseline_ft_weighted_wer"] is not None: summary["baseline_ft_vs_base_model_weighted_wer"] = ( summary["base_model_weighted_wer"] - summary["baseline_ft_weighted_wer"] ) else: summary["baseline_ft_vs_base_model_weighted_wer"] = None if summary["base_model_weighted_cer"] is not None and summary["baseline_ft_weighted_cer"] is not None: summary["baseline_ft_vs_base_model_weighted_cer"] = ( summary["base_model_weighted_cer"] - summary["baseline_ft_weighted_cer"] ) else: summary["baseline_ft_vs_base_model_weighted_cer"] = None if summary["base_model_weighted_wer"] is not None and summary["latent_reasoning_weighted_wer"] is not None: summary["latent_reasoning_vs_base_model_weighted_wer"] = ( summary["base_model_weighted_wer"] - summary["latent_reasoning_weighted_wer"] ) else: summary["latent_reasoning_vs_base_model_weighted_wer"] = None if summary["base_model_weighted_cer"] is not None and summary["latent_reasoning_weighted_cer"] is not None: summary["latent_reasoning_vs_base_model_weighted_cer"] = ( summary["base_model_weighted_cer"] - summary["latent_reasoning_weighted_cer"] ) else: summary["latent_reasoning_vs_base_model_weighted_cer"] = None if summary["baseline_ft_weighted_wer"] is not None and summary["latent_reasoning_weighted_wer"] is not None: summary["latent_reasoning_vs_baseline_ft_weighted_wer"] = ( summary["baseline_ft_weighted_wer"] - summary["latent_reasoning_weighted_wer"] ) else: summary["latent_reasoning_vs_baseline_ft_weighted_wer"] = None if summary["baseline_ft_weighted_cer"] is not None and summary["latent_reasoning_weighted_cer"] is not None: summary["latent_reasoning_vs_baseline_ft_weighted_cer"] = ( summary["baseline_ft_weighted_cer"] - summary["latent_reasoning_weighted_cer"] ) else: summary["latent_reasoning_vs_baseline_ft_weighted_cer"] = None if summary["base_model_weighted_wer"] is not None and summary["prompt_tuning_weighted_wer"] is not None: summary["prompt_tuning_vs_base_model_weighted_wer"] = ( summary["base_model_weighted_wer"] - summary["prompt_tuning_weighted_wer"] ) else: summary["prompt_tuning_vs_base_model_weighted_wer"] = None if summary["base_model_weighted_cer"] is not None and summary["prompt_tuning_weighted_cer"] is not None: summary["prompt_tuning_vs_base_model_weighted_cer"] = ( summary["base_model_weighted_cer"] - summary["prompt_tuning_weighted_cer"] ) else: summary["prompt_tuning_vs_base_model_weighted_cer"] = None if summary["base_model_weighted_wer"] is not None and summary["lora_r16_weighted_wer"] is not None: summary["lora_r16_vs_base_model_weighted_wer"] = ( summary["base_model_weighted_wer"] - summary["lora_r16_weighted_wer"] ) else: summary["lora_r16_vs_base_model_weighted_wer"] = None if summary["base_model_weighted_cer"] is not None and summary["lora_r16_weighted_cer"] is not None: summary["lora_r16_vs_base_model_weighted_cer"] = ( summary["base_model_weighted_cer"] - summary["lora_r16_weighted_cer"] ) else: summary["lora_r16_vs_base_model_weighted_cer"] = None if summary["prompt_tuning_weighted_wer"] is not None and summary["latent_reasoning_weighted_wer"] is not None: summary["latent_reasoning_vs_prompt_tuning_weighted_wer"] = ( summary["prompt_tuning_weighted_wer"] - summary["latent_reasoning_weighted_wer"] ) else: summary["latent_reasoning_vs_prompt_tuning_weighted_wer"] = None if summary["prompt_tuning_weighted_cer"] is not None and summary["latent_reasoning_weighted_cer"] is not None: summary["latent_reasoning_vs_prompt_tuning_weighted_cer"] = ( summary["prompt_tuning_weighted_cer"] - summary["latent_reasoning_weighted_cer"] ) else: summary["latent_reasoning_vs_prompt_tuning_weighted_cer"] = None if summary["lora_r16_weighted_wer"] is not None and summary["latent_reasoning_weighted_wer"] is not None: summary["latent_reasoning_vs_lora_r16_weighted_wer"] = ( summary["lora_r16_weighted_wer"] - summary["latent_reasoning_weighted_wer"] ) else: summary["latent_reasoning_vs_lora_r16_weighted_wer"] = None if summary["lora_r16_weighted_cer"] is not None and summary["latent_reasoning_weighted_cer"] is not None: summary["latent_reasoning_vs_lora_r16_weighted_cer"] = ( summary["lora_r16_weighted_cer"] - summary["latent_reasoning_weighted_cer"] ) else: summary["latent_reasoning_vs_lora_r16_weighted_cer"] = None print("\n=== Weighted Summary ===") for k, v in summary.items(): if v is None: print(f"{k}: -") else: print(f"{k}: {v:.6f}") if args.output_json: payload = { "model_id": args.model_id, "dataset_name": dataset_name, "dataset_name_input": args.dataset_name, "split": args.split, "configs": configs, "streaming": args.streaming, "text_normalizer": args.text_normalizer, "base_model_enabled": run_base_model, "baseline_checkpoint": args.baseline_ckpt, "prompt_tuning_checkpoint": args.prompt_tuning_ckpt, "lora_r16_checkpoint": args.lora_r16_ckpt, "latent_reasoning_checkpoint": args.latent_ckpt, "latent_checkpoint": args.latent_ckpt, "rows": merged_rows, "summary": summary, } with open(args.output_json, "w", encoding="utf-8") as f: json.dump(payload, f, ensure_ascii=False, indent=2) print(f"\nSaved JSON: {args.output_json}") if __name__ == "__main__": main()