scripts/run_public_checkpoint_comparison.py

#!/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 = {"<blank>", "<blk>", "<eps>", "<pad>", "[pad]"}
SPECIAL_LABELS = BLANK_LABELS | {"", "<bos>", "<eos>", "<unk>", "[unk]", "[UNK]"}
SILENCE_LABELS = {"sil", "[sil]", "<sil>", "!sil", "sp", "<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"<missing:{idx}>" 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"<missing:{idx}>") 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