#!/usr/bin/env python """Run comparable log-prob, forced-alignment, and text-free timestamp exports. The runner is intentionally registry-driven: adding a checkpoint should usually mean adding one JSON object to the registry, not copying this script. Example: python scripts/run_public_checkpoint_comparison.py \ --registry scripts/public_four_model_registry.json \ --manifest /path/to/test.json \ --reference-field canonical_aligned \ --output-dir public_checkpoint_runs/demo \ --device cuda \ --limit 8 \ --overwrite """ from __future__ import annotations import argparse import csv import importlib.util import json import math import re import sys from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterable import numpy as np import torch REPO_ROOT = Path(__file__).resolve().parents[1] DEFAULT_REGISTRY = REPO_ROOT / "scripts" / "public_four_model_registry.json" DEFAULT_OUTPUT_DIR = REPO_ROOT / "public_checkpoint_runs" / "four_model_compare" FALLBACK_ENCODER_ASR = REPO_ROOT / "pretrained_models" / "CTC_for_IF-MDD" / "MyEncoderASR.py" if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) BLANK_LABELS = {"", "", "", "", "[pad]"} SPECIAL_LABELS = BLANK_LABELS | {"", "", "", "", "[unk]", "[UNK]"} SILENCE_LABELS = {"sil", "[sil]", "", "!sil", "sp", "", "spn", "pau", "h#"} TOPOLOGY_SUFFIX_RE = re.compile(r"_(?:state)?\d+$") @dataclass(frozen=True) class ModelSpec: model_id: str display: str backend: str source: str source_is_local: bool hparams_file: str model_kind: str sample_rate: int encoder_py: Path | None variant: str | None = None @dataclass(frozen=True) class Utterance: utt_id: str wav: Path duration: float metadata: dict[str, Any] @dataclass class ModelRunSummary: model_id: str display: str model_kind: str source: str log_probs_h5: str text_free_dir: str forced_alignment_dir: str utterances: int text_free_segments: int forced_alignment_segments: int forced_alignment_failures: int skipped_forced_alignment: int def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--registry", type=Path, default=DEFAULT_REGISTRY) parser.add_argument( "--model-filter", default="", help="Optional regex over registry model ids/display names, useful for smoke tests.", ) input_group = parser.add_mutually_exclusive_group(required=True) input_group.add_argument("--manifest", type=Path) input_group.add_argument("--wav", type=Path) parser.add_argument("--utt-id", default="", help="Utterance id for --wav input.") parser.add_argument( "--dataset-name", default="custom", help="Dataset name stored in JSONL summaries.", ) parser.add_argument( "--dataset-split", default="", help="Split name stored in outputs. Defaults to manifest stem or 'single'.", ) parser.add_argument( "--reference-field", default="canonical_aligned", help="Manifest field containing reference phones for forced alignment.", ) parser.add_argument( "--reference-phones", default="", help="Reference phone string for --wav input, e.g. 'sil dh ah ...'.", ) parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR) parser.add_argument("--device", default="cuda") parser.add_argument("--batch-size", type=int, default=1) parser.add_argument("--limit", type=int, default=0, help="0 means all utterances.") parser.add_argument("--overwrite", action="store_true") parser.add_argument("--dry-run", action="store_true") parser.add_argument( "--tasks", default="log_probs,forced_alignment,text_free", help=( "Comma-separated tasks. log_probs is always computed internally; " "valid values: log_probs, forced_alignment, text_free." ), ) parser.add_argument("--drop-silence", action="store_true") parser.add_argument("--preserve-case", action="store_true") parser.add_argument( "--blank-policy", choices=["drop", "previous", "next", "split"], default="previous", help="How to assign blank gaps around forced-alignment phone spans.", ) parser.add_argument("--progress-every", type=int, default=25) parser.add_argument("--decimals", type=int, default=6) return parser.parse_args() def safe_id(value: str) -> str: text = re.sub(r"[^A-Za-z0-9_.-]+", "_", str(value)).strip("._") return text or "model" def read_registry(path: Path) -> list[ModelSpec]: payload = json.loads(path.read_text(encoding="utf-8")) raw_models = payload.get("models", []) if isinstance(raw_models, dict): raw_models = [ {"id": key, **value} for key, value in raw_models.items() if isinstance(value, dict) ] if not isinstance(raw_models, list) or not raw_models: raise ValueError(f"{path} must contain a nonempty 'models' list or object.") models = [] for raw in raw_models: if not isinstance(raw, dict): raise ValueError(f"Bad model entry in {path}: {raw!r}") model_id = str(raw.get("id") or raw.get("name") or "").strip() if not model_id: raise ValueError(f"Model entry is missing id: {raw!r}") source_text = str(raw.get("source") or "").strip() if not source_text: raise ValueError(f"Model entry {model_id!r} is missing source.") source, source_is_local = resolve_model_source(source_text, path.parent) encoder_py = resolve_optional_path(raw.get("encoder_py"), path.parent) models.append( ModelSpec( model_id=safe_id(model_id), display=str(raw.get("display") or model_id), backend=str(raw.get("backend") or "speechbrain"), source=source, source_is_local=source_is_local, hparams_file=str(raw.get("hparams_file") or "inference.yaml"), model_kind=str(raw.get("model_kind") or raw.get("kind") or ""), sample_rate=int(raw.get("sample_rate") or 16000), encoder_py=encoder_py, variant=str(raw.get("variant") or "") or None, ) ) return models def resolve_model_source(value: str, registry_dir: Path) -> tuple[str, bool]: path = Path(value).expanduser() candidates = [] if path.is_absolute(): candidates.append(path) else: candidates.extend([registry_dir / path, REPO_ROOT / path, Path.cwd() / path]) for candidate in candidates: if candidate.exists(): return str(candidate.resolve()), True return value, False def resolve_optional_path(value: object, registry_dir: Path) -> Path | None: if not value: return None text = str(value) path = Path(text).expanduser() candidates = [path] if path.is_absolute() else [registry_dir / path, REPO_ROOT / path, Path.cwd() / path] for candidate in candidates: if candidate.exists(): return candidate.resolve() return candidates[0] def load_utterances(args: argparse.Namespace) -> list[Utterance]: if args.wav is not None: wav = args.wav.expanduser().resolve() utt_id = args.utt_id or wav.stem metadata: dict[str, Any] = {} if args.reference_phones: metadata[args.reference_field] = args.reference_phones return [Utterance(utt_id=utt_id, wav=wav, duration=0.0, metadata=metadata)] assert args.manifest is not None data = json.loads(args.manifest.read_text(encoding="utf-8")) raw_items: Iterable[tuple[object, object]] if isinstance(data, dict): raw_items = data.items() elif isinstance(data, list): raw_items = enumerate(data) else: raise ValueError(f"Unsupported manifest payload: {type(data).__name__}") utterances = [] for key, value in raw_items: if not isinstance(value, dict): continue wav_text = value.get("wav") or value.get("wav_path") or value.get("audio") if not wav_text: continue wav = resolve_wav_path(str(wav_text), args.manifest.parent) utt_id = str(value.get("id") or value.get("utt_id") or value.get("utt") or Path(str(wav_text)).stem or key) duration = safe_float(value.get("duration"), default=0.0) metadata = dict(value) metadata["_manifest_key"] = str(key) metadata["_resolved_wav"] = str(wav) utterances.append(Utterance(utt_id=utt_id, wav=wav, duration=duration, metadata=metadata)) if args.limit > 0 and len(utterances) >= args.limit: break if not utterances: raise ValueError(f"No manifest entries with wav/audio paths found in {args.manifest}") return utterances def resolve_wav_path(text: str, base_dir: Path) -> Path: raw = Path(text).expanduser() candidates = [] if raw.is_absolute(): candidates.append(raw) else: candidates.extend([base_dir / raw, REPO_ROOT / raw, Path.cwd() / raw]) replacements = [ ("/home/m64000/work/dataset/", "/work/gm64/m64000/dataset/"), ("/home/m64000/work/IF-MDD/", "/work/gm64/m64000/IF-MDD/"), ( "/home/kevingenghaopeng/MDD/IF-MDD/data/speechocean762/wav/", "/work/gm64/m64000/dataset/speechocean762_flat/speechocean762/wav/", ), ] for old, new in replacements: if text.startswith(old): candidates.append(Path(text.replace(old, new, 1))) if text.endswith(".WAV"): candidates.append(Path(text[:-4] + ".wav")) elif text.endswith(".wav"): candidates.append(Path(text[:-4] + ".WAV")) stem = raw.stem candidates.extend( [ Path("/work/gm64/m64000/dataset/speechocean762_flat/speechocean762/wav") / f"{stem}.wav", Path("/work/gm64/m64000/dataset/speechocean762_flat/speechocean762/wav") / f"{stem}.WAV", ] ) for candidate in candidates: if candidate.exists(): return candidate.resolve() return candidates[0] def safe_float(value: object, default: float = 0.0) -> float: try: out = float(value) except Exception: return default return out if math.isfinite(out) else default def import_module_from_path(path: Path, module_name: str): spec = importlib.util.spec_from_file_location(module_name, path) if spec is None or spec.loader is None: raise RuntimeError(f"Could not import module from {path}") module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module spec.loader.exec_module(module) return module class SpeechBrainRunner: def __init__(self, spec: ModelSpec, device: str): if spec.backend != "speechbrain": raise ValueError(f"Unsupported backend for this runner: {spec.backend}") self.spec = spec self.device = device self.model = None def load(self): module_path = self.spec.encoder_py source_path = Path(self.spec.source) if self.spec.source_is_local else None if module_path is None and source_path is not None: for filename in ["TopologyEncoderASR.py", "MyEncoderASR.py"]: candidate = source_path / filename if candidate.exists(): module_path = candidate break if module_path is None and FALLBACK_ENCODER_ASR.exists(): module_path = FALLBACK_ENCODER_ASR if module_path is not None and module_path.exists(): module = import_module_from_path(module_path, f"public_runner_{self.spec.model_id}") cls = getattr(module, "TopologyEncoderASR", None) or getattr(module, "MyEncoderASR", None) if cls is None: raise RuntimeError(f"{module_path} does not define TopologyEncoderASR or MyEncoderASR") else: try: from speechbrain.inference.ASR import EncoderASR except Exception as exc: raise RuntimeError( "speechbrain is required when no local MyEncoderASR.py is available." ) from exc cls = EncoderASR # SpeechBrain evaluates relative paths in HyperPyYAML against the # process working directory. Public Space bundles are mounted under # /data, so force all bundle-local paths to the resolved model source. overrides = {} if self.spec.source_is_local: bundle_dir = str(Path(self.spec.source).resolve()) overrides = { "save_folder": bundle_dir, "output_folder": bundle_dir, "pretrained_models_path": bundle_dir, "lab_enc_file": str(Path(bundle_dir) / "label_encoder.txt"), } self.model = cls.from_hparams( source=self.spec.source, hparams_file=self.spec.hparams_file, run_opts={"device": self.device}, overrides=overrides, overrides_must_match=False, ) if hasattr(self.model, "mods"): self.model.mods.eval() return self def load_audio_batch(self, batch: list[Utterance]) -> tuple[torch.Tensor, torch.Tensor, list[int]]: if self.model is None: raise RuntimeError("Model is not loaded.") waves, sample_lengths = [], [] for item in batch: if not item.wav.exists(): raise FileNotFoundError(f"Missing wav for {item.utt_id}: {item.wav}") wav = self.model.load_audio(str(item.wav)) if wav.dim() > 1: wav = wav.squeeze() wav = wav.float() waves.append(wav) sample_lengths.append(int(wav.numel())) max_len = max(sample_lengths) wavs = torch.zeros(len(waves), max_len, dtype=torch.float32) for idx, wav in enumerate(waves): wavs[idx, : wav.numel()] = wav rel_lens = torch.tensor([length / max_len for length in sample_lengths], dtype=torch.float32) return wavs, rel_lens, sample_lengths @torch.inference_mode() def encode_batch(self, batch: list[Utterance]) -> tuple[torch.Tensor, list[int], list[int]]: wavs, rel_lens, sample_lengths = self.load_audio_batch(batch) encoded = self.model.encode_batch(wavs, rel_lens).detach().float().cpu() frame_lengths = (rel_lens * encoded.shape[1]).long().clamp(min=1, max=encoded.shape[1]).tolist() return encoded, sample_lengths, frame_lengths def labels(self) -> list[str]: if self.model is None: raise RuntimeError("Model is not loaded.") labels = labels_from_tokenizer(getattr(self.model, "tokenizer", None)) if labels: return labels source_path = Path(self.spec.source) if self.spec.source_is_local else None if source_path is not None: for label_file in (source_path / "label_encoder.txt", source_path / "content_label_encoder.txt", source_path / "subphonetic_topology" / "topology_label_encoder.txt"): if label_file.exists(): labels = parse_speechbrain_label_encoder(label_file) if labels: return labels raise RuntimeError(f"Could not recover labels for {self.spec.model_id}") class FrontendRunner: """Adapter for the standalone LibriSpeech acoustic frontend exports.""" def __init__(self, spec: ModelSpec, device: str): self.spec = spec self.device = device self.model = None self._labels: list[str] = [] def load(self): if not self.spec.source_is_local: raise ValueError("The standalone frontend backend requires a local mounted bundle.") source = Path(self.spec.source) module_path = source / "modeling_ifmdd_acoustic_frontend.py" module = import_module_from_path(module_path, f"public_frontend_{self.spec.model_id}") cls = getattr(module, "IFMDDAcousticFrontend") self.model = cls.from_pretrained( source, variant=self.spec.variant or "20ms", map_location=self.device, local_files_only=True, ) self.model.to(self.device) self.model.eval() label_path = source / "labels" / "topology_label_encoder.txt" self._labels = parse_speechbrain_label_encoder(label_path) return self def labels(self) -> list[str]: return self._labels def load_audio_batch(self, batch: list[Utterance]) -> tuple[torch.Tensor, torch.Tensor, list[int]]: waves, sample_lengths = [], [] for item in batch: if not item.wav.exists(): raise FileNotFoundError(f"Missing wav for {item.utt_id}") import soundfile as sf import torchaudio import torch.nn.functional as F samples, sample_rate = sf.read(str(item.wav), always_2d=False) wav = torch.as_tensor(samples, dtype=torch.float32) if wav.ndim > 1: wav = wav.mean(dim=-1) if int(sample_rate) != self.spec.sample_rate: wav = torchaudio.functional.resample(wav, int(sample_rate), self.spec.sample_rate) waves.append(wav) sample_lengths.append(int(wav.numel())) max_len = max(sample_lengths) padded = torch.zeros(len(waves), max_len, dtype=torch.float32) for index, wav in enumerate(waves): padded[index, : wav.numel()] = wav rel_lens = torch.tensor([length / max_len for length in sample_lengths], dtype=torch.float32) return padded.to(self.device), rel_lens.to(self.device), sample_lengths @torch.inference_mode() def encode_batch(self, waves: torch.Tensor, rel_lens: torch.Tensor) -> torch.Tensor: outputs = self.model.extract(waves, input_lengths=(rel_lens * waves.shape[1]).long()) # General ours bundles expose topology states; they are intentionally # registered only as Subphonetic OTTC. return outputs["fused_topology_log_probs"] def labels_from_tokenizer(tokenizer: object | None) -> list[str]: if tokenizer is None: return [] ind2lab = getattr(tokenizer, "ind2lab", None) if isinstance(ind2lab, dict) and ind2lab: return [str(ind2lab[idx]) for idx in sorted(ind2lab)] if isinstance(ind2lab, (list, tuple)) and ind2lab: return [str(item) for item in ind2lab] lab2ind = getattr(tokenizer, "lab2ind", None) if isinstance(lab2ind, dict) and lab2ind: size = max(int(idx) for idx in lab2ind.values()) + 1 labels = [f"" for idx in range(size)] for label, idx in lab2ind.items(): labels[int(idx)] = str(label) return labels return [] def parse_speechbrain_label_encoder(path: Path) -> list[str]: labels: dict[int, str] = {} for line in path.read_text(encoding="utf-8").splitlines(): stripped = line.strip() if not stripped or stripped.startswith("=") or "=>" not in stripped: continue label, index = [part.strip() for part in stripped.split("=>", 1)] label = label.strip("'\"") try: labels[int(index)] = label except ValueError: continue if not labels: return [] return [labels.get(idx, f"") for idx in range(max(labels) + 1)] def choose_blank_id(labels: list[str]) -> int: for idx, label in enumerate(labels): if str(label).strip().lower() in BLANK_LABELS: return int(idx) return 0 def normalize_label(label: str, preserve_case: bool, drop_silence: bool) -> str | None: text = str(label).strip() if not preserve_case: text = text.lower() text = TOPOLOGY_SUFFIX_RE.sub("", text) if text.lower() in SPECIAL_LABELS: return None if drop_silence and text.lower() in SILENCE_LABELS: return None if text == "[sil]": return "sil" return text def frame_times( start_frame: int, end_frame: int, frame_len: int, duration: float, ) -> tuple[float, float]: if duration <= 0.0: return float(start_frame), float(end_frame) start = float(start_frame) * duration / float(max(1, frame_len)) end = float(end_frame) * duration / float(max(1, frame_len)) return min(max(start, 0.0), duration), min(max(end, start), duration) def confidence(log_values: Iterable[float]) -> float: values = np.asarray(list(log_values), dtype=np.float64) if values.size == 0: return 0.0 return float(np.exp(np.clip(values, -80.0, 0.0)).mean()) def posterior_best_path_segments( log_probs: torch.Tensor, labels: list[str], blank_id: int, duration: float, preserve_case: bool, drop_silence: bool, ) -> list[dict[str, Any]]: matrix = log_probs.detach().cpu().float() frame_len = int(matrix.shape[0]) frame_ids = matrix.argmax(dim=-1).tolist() frame_scores = matrix[torch.arange(frame_len), torch.tensor(frame_ids)].tolist() segments: list[dict[str, Any]] = [] current_label: str | None = None start_frame = 0 score_values: list[float] = [] def flush(end_frame: int) -> None: nonlocal current_label, start_frame, score_values if current_label is None: score_values = [] return start, end = frame_times(start_frame, end_frame, frame_len, duration) if end > start: segments.append( { "phone": current_label, "onset": start, "offset": end, "start_frame": int(start_frame), "end_frame": int(end_frame), "score": confidence(score_values), "log_probability": float(np.mean(score_values)) if score_values else float("-inf"), } ) current_label = None score_values = [] for frame_idx, (token_id, score) in enumerate(zip(frame_ids, frame_scores)): raw_label = labels[int(token_id)] if 0 <= int(token_id) < len(labels) else "" label = normalize_label(raw_label, preserve_case=preserve_case, drop_silence=drop_silence) if int(token_id) == blank_id: label = None if label is not None and label == current_label: score_values.append(float(score)) continue flush(frame_idx) if label is None: start_frame = frame_idx + 1 continue current_label = label start_frame = frame_idx score_values = [float(score)] flush(frame_len) return segments def phones_from_value(value: Any, preserve_case: bool) -> list[str]: if value is None: return [] raw = value if isinstance(value, list) else str(value).split() phones = [] for phone in raw: text = str(phone).strip() if not text: continue phones.append(text if preserve_case else text.lower()) return phones def label_lookup(labels: list[str]) -> dict[str, int]: lookup: dict[str, int] = {} for idx, label in enumerate(labels): text = str(label) lookup.setdefault(text, int(idx)) lookup.setdefault(text.lower(), int(idx)) return lookup def split_subphone_label(label: str) -> tuple[str, int | None]: match = re.search(r"_(?:state)?(\d+)$", str(label)) if not match: return str(label), None return str(label)[: match.start()], int(match.group(1)) def topology_state_inventory(labels: list[str]) -> dict[str, list[tuple[int, int]]]: inventory: dict[str, list[tuple[int, int]]] = {} for token_id, label in enumerate(labels): text = str(label) if text.lower() in SPECIAL_LABELS: continue phone, state = split_subphone_label(text) if state is None: continue inventory.setdefault(phone.lower(), []).append((state, int(token_id))) return {phone: sorted(items, key=lambda item: item[0]) for phone, items in inventory.items()} def prepare_forced_alignment_target( phones: list[str], labels: list[str], ) -> tuple[list[int], list[int], str]: lookup = label_lookup(labels) direct_ids = [] missing = [] for phone in phones: if phone in lookup: direct_ids.append(lookup[phone]) elif phone.lower() in lookup: direct_ids.append(lookup[phone.lower()]) else: missing.append(phone) if not missing: return direct_ids, [1] * len(phones), "content" inventory = topology_state_inventory(labels) token_ids = [] group_sizes = [] topology_missing = [] for phone in phones: states = inventory.get(phone.lower()) if not states: topology_missing.append(phone) continue group_sizes.append(len(states)) token_ids.extend(token_id for _state, token_id in states) if topology_missing: raise ValueError( "reference phones missing from model label inventory: " + ", ".join(sorted(set(topology_missing))[:20]) ) return token_ids, group_sizes, "topology" def ctc_min_frames(token_ids: list[int]) -> int: repeats = sum(1 for prev, cur in zip(token_ids, token_ids[1:]) if prev == cur) return len(token_ids) + repeats def apply_blank_policy(segments: list[dict[str, Any]], frame_len: int, policy: str) -> None: if policy == "drop" or not segments: return original_starts = [int(segment["start_frame"]) for segment in segments] original_ends = [int(segment["end_frame"]) for segment in segments] segments[0]["start_frame"] = 0 if policy == "previous": for idx in range(len(segments) - 1): segments[idx]["end_frame"] = original_starts[idx + 1] segments[-1]["end_frame"] = frame_len elif policy == "next": for idx in range(1, len(segments)): segments[idx]["start_frame"] = original_ends[idx - 1] segments[-1]["end_frame"] = frame_len elif policy == "split": for idx in range(len(segments) - 1): midpoint = int(round((original_ends[idx] + original_starts[idx + 1]) / 2.0)) segments[idx]["end_frame"] = midpoint segments[idx + 1]["start_frame"] = midpoint segments[-1]["end_frame"] = frame_len else: raise ValueError(f"Unknown blank policy: {policy}") def force_align_segments( log_probs: torch.Tensor, reference_phones: list[str], labels: list[str], blank_id: int, duration: float, blank_policy: str, ) -> tuple[list[dict[str, Any]], str]: try: from torchaudio.functional import forced_align, merge_tokens except Exception as exc: raise RuntimeError("torchaudio.functional.forced_align is required for FA") from exc matrix = log_probs.detach().cpu().float() frame_len, _classes = matrix.shape token_ids, group_sizes, target_mode = prepare_forced_alignment_target(reference_phones, labels) min_frames = ctc_min_frames(token_ids) if frame_len < min_frames: raise ValueError( f"too few frames for CTC FA: input_frames={frame_len} min_frames={min_frames}" ) alignments, scores = forced_align( log_probs=matrix.unsqueeze(0), targets=torch.tensor([token_ids], dtype=torch.int32), input_lengths=torch.tensor([frame_len], dtype=torch.int32), target_lengths=torch.tensor([len(token_ids)], dtype=torch.int32), blank=int(blank_id), ) spans = merge_tokens(alignments[0], scores[0].exp()) state_segments = [ { "token_id": int(span.token), "start_frame": int(span.start), "end_frame": int(span.end), "score": float(span.score), } for span in spans ] if len(state_segments) != len(token_ids): raise ValueError( f"FA recovered {len(state_segments)} spans for {len(token_ids)} target tokens" ) apply_blank_policy(state_segments, frame_len=frame_len, policy=blank_policy) for segment in state_segments: start, end = frame_times( int(segment["start_frame"]), int(segment["end_frame"]), frame_len, duration, ) segment["onset"] = start segment["offset"] = end if target_mode == "content": for segment, phone in zip(state_segments, reference_phones): segment["phone"] = phone return state_segments, target_mode collapsed = [] cursor = 0 for phone, size in zip(reference_phones, group_sizes): group = state_segments[cursor : cursor + size] cursor += size collapsed.append( { "phone": phone, "token_id": int(group[0]["token_id"]), "start_frame": int(group[0]["start_frame"]), "end_frame": int(group[-1]["end_frame"]), "onset": float(group[0]["onset"]), "offset": float(group[-1]["offset"]), "score": float(sum(float(item["score"]) for item in group) / max(1, len(group))), } ) return collapsed, target_mode def ctm_line(utt: str, segment: dict[str, Any], decimals: int, include_score: bool) -> str: onset = float(segment["onset"]) duration = max(0.0, float(segment["offset"]) - onset) line = f"{utt} 1 {onset:.{decimals}f} {duration:.{decimals}f} {segment['phone']}" if include_score: line += f" {float(segment.get('score', 0.0)):.{decimals}f}" return line def write_segments_tsv_header(path: Path) -> Any: f = path.open("w", encoding="utf-8", newline="") writer = csv.writer(f, delimiter="\t") writer.writerow( [ "model", "utt", "wav", "phone", "onset", "offset", "duration", "score", "start_frame", "end_frame", ] ) return f, writer def write_alignment_record( handle: Any, *, dataset: str, split: str, alignment_type: str, model: ModelSpec, utt: Utterance, duration: float, segments: list[dict[str, Any]], extra: dict[str, Any] | None = None, ) -> None: record = { "dataset": dataset, "dataset_split": split, "alignment_type": alignment_type, "model": model.model_id, "model_display": model.display, "model_kind": model.model_kind, "utt": utt.utt_id, "wav": str(utt.wav), "duration": duration, "phones": [segment["phone"] for segment in segments], "pred_starts": [float(segment["onset"]) for segment in segments], "pred_ends": [float(segment["offset"]) for segment in segments], "mean_scores": [float(segment.get("score", 0.0)) for segment in segments], "start_frames": [int(segment["start_frame"]) for segment in segments], "end_frames": [int(segment["end_frame"]) for segment in segments], "segments": segments, } if extra: record.update(extra) handle.write(json.dumps(record, ensure_ascii=False) + "\n") def write_segment_rows( writer: csv.writer, model: ModelSpec, utt: Utterance, duration: float, segments: list[dict[str, Any]], decimals: int, ) -> None: for segment in segments: onset = float(segment["onset"]) offset = float(segment["offset"]) writer.writerow( [ model.model_id, utt.utt_id, str(utt.wav), segment["phone"], f"{onset:.{decimals}f}", f"{offset:.{decimals}f}", f"{max(0.0, offset - onset):.{decimals}f}", f"{float(segment.get('score', 0.0)):.{decimals}f}", int(segment["start_frame"]), int(segment["end_frame"]), ] ) def h5_string_dtype(): import h5py return h5py.string_dtype(encoding="utf-8") def open_log_h5(path: Path, model: ModelSpec, labels: list[str], utterances: list[Utterance], args: argparse.Namespace): import h5py if path.exists() and not args.overwrite: raise FileExistsError(f"{path} exists; pass --overwrite") string_dtype = h5_string_dtype() h5 = h5py.File(path, "w") h5.attrs["created_utc"] = datetime.now(timezone.utc).isoformat() h5.attrs["command"] = " ".join(sys.argv) h5.attrs["model_id"] = model.model_id h5.attrs["model_display"] = model.display h5.attrs["model_kind"] = model.model_kind h5.attrs["model_source"] = model.source h5.attrs["backend"] = model.backend h5.attrs["value_kind"] = "log_probs" h5.attrs["tensor_semantics"] = "frame-level CTC log probabilities from model.encode_batch" h5.attrs["padded"] = True h5.attrs["pad_value"] = -1.0e30 h5.create_dataset("ids", data=[u.utt_id for u in utterances], dtype=string_dtype) h5.create_dataset("wav_paths", data=[str(u.wav) for u in utterances], dtype=string_dtype) h5.create_dataset( "utterance_metadata_json", data=[json.dumps(u.metadata, ensure_ascii=False) for u in utterances], dtype=string_dtype, ) h5.create_dataset("labels", data=labels, dtype=string_dtype) h5.create_dataset("durations", shape=(len(utterances),), dtype=np.float32) h5.create_dataset("sample_lengths", shape=(len(utterances),), dtype=np.int32) h5.create_dataset("frame_lengths", shape=(len(utterances),), dtype=np.int32) return h5 def iter_batches(items: list[Utterance], batch_size: int): for start in range(0, len(items), batch_size): yield start, items[start : start + batch_size] def selected_tasks(value: str) -> set[str]: tasks = {item.strip() for item in value.split(",") if item.strip()} valid = {"log_probs", "forced_alignment", "text_free"} unknown = tasks - valid if unknown: raise ValueError(f"Unknown tasks: {', '.join(sorted(unknown))}") tasks.add("log_probs") return tasks def run_model( spec: ModelSpec, utterances: list[Utterance], args: argparse.Namespace, tasks: set[str], ) -> ModelRunSummary: out_dir = args.output_dir / spec.model_id text_free_dir = out_dir / "text_free" fa_dir = out_dir / "forced_alignment" out_dir.mkdir(parents=True, exist_ok=True) text_free_dir.mkdir(parents=True, exist_ok=True) fa_dir.mkdir(parents=True, exist_ok=True) runner = SpeechBrainRunner(spec, args.device).load() labels = runner.labels() blank_id = choose_blank_id(labels) h5_path = out_dir / "log_probs.h5" h5 = open_log_h5(h5_path, spec, labels, utterances, args) text_free_align = (text_free_dir / "alignments.jsonl").open("w", encoding="utf-8") text_free_ctm = (text_free_dir / "phone.ctm").open("w", encoding="utf-8") text_free_score_ctm = (text_free_dir / "phone.score.ctm").open("w", encoding="utf-8") text_free_seg_f, text_free_seg_writer = write_segments_tsv_header(text_free_dir / "segments.tsv") fa_align = (fa_dir / "alignments.jsonl").open("w", encoding="utf-8") fa_ctm = (fa_dir / "phone.ctm").open("w", encoding="utf-8") fa_score_ctm = (fa_dir / "phone.score.ctm").open("w", encoding="utf-8") fa_seg_f, fa_seg_writer = write_segments_tsv_header(fa_dir / "segments.tsv") fa_errors: list[dict[str, Any]] = [] log_dset = None max_frames = 0 num_classes = None text_free_segments = 0 fa_segments = 0 skipped_fa = 0 split = args.dataset_split or (args.manifest.stem if args.manifest else "single") try: for batch_start, batch in iter_batches(utterances, args.batch_size): encoded, sample_lengths, frame_lengths = runner.encode_batch(batch) batch_size, batch_frames, classes = encoded.shape if num_classes is None: num_classes = int(classes) log_dset = h5.create_dataset( "log_probs", shape=(len(utterances), 0, classes), maxshape=(len(utterances), None, classes), chunks=(1, min(256, max(1, batch_frames)), classes), dtype=np.float32, fillvalue=-1.0e30, ) if log_dset is None: raise RuntimeError("log_probs dataset was not initialized") if int(classes) != int(num_classes): raise RuntimeError(f"class count changed for {spec.model_id}: {num_classes} -> {classes}") if batch_frames > max_frames: log_dset.resize((len(utterances), batch_frames, classes)) max_frames = int(batch_frames) for local_idx, utt in enumerate(batch): row = batch_start + local_idx frame_len = int(frame_lengths[local_idx]) sample_len = int(sample_lengths[local_idx]) duration = float(utt.duration) if utt.duration > 0.0 else sample_len / float(spec.sample_rate) matrix = encoded[local_idx, :frame_len, :].contiguous() log_dset[row, :frame_len, :] = matrix.numpy() h5["durations"][row] = duration h5["sample_lengths"][row] = sample_len h5["frame_lengths"][row] = frame_len if "text_free" in tasks: segments = posterior_best_path_segments( matrix, labels=labels, blank_id=blank_id, duration=duration, preserve_case=args.preserve_case, drop_silence=args.drop_silence, ) write_alignment_record( text_free_align, dataset=args.dataset_name, split=split, alignment_type="posterior_best_path", model=spec, utt=utt, duration=duration, segments=segments, extra={"blank_id": blank_id, "blank_label": labels[blank_id]}, ) write_segment_rows(text_free_seg_writer, spec, utt, duration, segments, args.decimals) for segment in segments: text_free_ctm.write(ctm_line(utt.utt_id, segment, args.decimals, include_score=False) + "\n") text_free_score_ctm.write(ctm_line(utt.utt_id, segment, args.decimals, include_score=True) + "\n") text_free_segments += len(segments) if "forced_alignment" in tasks: reference_phones = phones_from_value( utt.metadata.get(args.reference_field), preserve_case=args.preserve_case, ) if not reference_phones: skipped_fa += 1 fa_errors.append( { "utt": utt.utt_id, "stage": "reference", "error": f"missing reference field {args.reference_field!r}", } ) else: try: segments, target_mode = force_align_segments( matrix, reference_phones=reference_phones, labels=labels, blank_id=blank_id, duration=duration, blank_policy=args.blank_policy, ) write_alignment_record( fa_align, dataset=args.dataset_name, split=split, alignment_type=f"{args.reference_field}_fa", model=spec, utt=utt, duration=duration, segments=segments, extra={ "target": args.reference_field, args.reference_field: " ".join(reference_phones), "fa_backend": "torchaudio", "target_mode": target_mode, "blank_id": blank_id, "blank_label": labels[blank_id], "blank_policy": args.blank_policy, }, ) write_segment_rows(fa_seg_writer, spec, utt, duration, segments, args.decimals) for segment in segments: fa_ctm.write(ctm_line(utt.utt_id, segment, args.decimals, include_score=False) + "\n") fa_score_ctm.write(ctm_line(utt.utt_id, segment, args.decimals, include_score=True) + "\n") fa_segments += len(segments) except Exception as exc: fa_errors.append( { "utt": utt.utt_id, "stage": "forced_alignment", "error": f"{type(exc).__name__}: {exc}", } ) processed = batch_start + batch_size if args.progress_every and ( processed >= len(utterances) or processed % args.progress_every == 0 or batch_start == 0 ): print( f"{spec.model_id}: {min(processed, len(utterances))}/{len(utterances)} " f"max_frames={max_frames} classes={classes}" ) finally: h5.attrs["num_utterances"] = len(utterances) h5.attrs["max_frames"] = int(max_frames) h5.attrs["num_classes"] = int(num_classes or 0) h5.attrs["blank_id"] = int(blank_id) h5.attrs["blank_label"] = labels[blank_id] if 0 <= blank_id < len(labels) else "" h5.close() text_free_align.close() text_free_ctm.close() text_free_score_ctm.close() text_free_seg_f.close() fa_align.close() fa_ctm.close() fa_score_ctm.close() fa_seg_f.close() write_errors_csv(fa_dir / "alignment_errors.csv", fa_errors) summary = ModelRunSummary( model_id=spec.model_id, display=spec.display, model_kind=spec.model_kind, source=spec.source, log_probs_h5=str(h5_path), text_free_dir=str(text_free_dir), forced_alignment_dir=str(fa_dir), utterances=len(utterances), text_free_segments=text_free_segments, forced_alignment_segments=fa_segments, forced_alignment_failures=sum(1 for row in fa_errors if row.get("stage") == "forced_alignment"), skipped_forced_alignment=skipped_fa, ) write_model_summary(out_dir / "summary.json", summary) return summary def write_errors_csv(path: Path, rows: list[dict[str, Any]]) -> None: with path.open("w", encoding="utf-8", newline="") as f: writer = csv.DictWriter(f, fieldnames=["utt", "stage", "error"]) writer.writeheader() for row in rows: writer.writerow({key: row.get(key, "") for key in ["utt", "stage", "error"]}) def write_model_summary(path: Path, summary: ModelRunSummary) -> None: path.write_text( json.dumps(summary.__dict__, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) def main() -> None: args = parse_args() tasks = selected_tasks(args.tasks) models = read_registry(args.registry) if args.model_filter: regex = re.compile(args.model_filter) models = [ model for model in models if regex.search(model.model_id) or regex.search(model.display) ] if not models: raise ValueError(f"--model-filter matched no registry entries: {args.model_filter}") utterances = load_utterances(args) args.output_dir.mkdir(parents=True, exist_ok=True) if args.dry_run: print(f"registry={args.registry}") print(f"models={len(models)}") for model in models: print(f" {model.model_id}: {model.source}") print(f"utterances={len(utterances)}") for utt in utterances[:5]: print(f" {utt.utt_id}: {utt.wav}") return summaries = [] for spec in models: print(f"\n=== {spec.display} ({spec.model_id}) ===") summaries.append(run_model(spec, utterances, args, tasks)) combined = { "created_utc": datetime.now(timezone.utc).isoformat(), "registry": str(args.registry), "dataset": args.dataset_name, "dataset_split": args.dataset_split or (args.manifest.stem if args.manifest else "single"), "reference_field": args.reference_field, "tasks": sorted(tasks), "utterances": len(utterances), "models": [summary.__dict__ for summary in summaries], } summary_path = args.output_dir / "summary.json" summary_path.write_text( json.dumps(combined, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) print(f"\nwrote combined summary: {summary_path}") if __name__ == "__main__": try: main() except Exception as exc: raise SystemExit(f"{type(exc).__name__}: {exc}") from exc