from __future__ import annotations import difflib import re from dataclasses import dataclass from faster_whisper import WhisperModel _STRIP_RE = re.compile(r"^[\s\.,!?;:\"'„“”‚‘’\(\)\[\]…—–-]+|[\s\.,!?;:\"'„“”‚‘’\(\)\[\]…—–-]+$") @dataclass class Word: text: str # normalized form for TTS (lowercased, punctuation stripped) raw: str # original Whisper token (kept for debugging / display) start: float # seconds in original audio end: float probability: float = float("nan") # Whisper per-word probability (NaN if unknown) @property def duration(self) -> float: return self.end - self.start def _normalize(token: str) -> str: return _STRIP_RE.sub("", token).lower() def transcribe_words( audio_path: str, model: WhisperModel, language: str = "cs", ) -> tuple[list[Word], str]: """Transcribe `audio_path` and return word-level segmentation + full text. Empty / pure-punctuation tokens are skipped. """ segments, _info = model.transcribe( audio_path, language=language, word_timestamps=True ) words: list[Word] = [] full_text: list[str] = [] for seg in segments: full_text.append(seg.text) if not seg.words: continue for w in seg.words: normalized = _normalize(w.word) if not normalized: continue prob = getattr(w, "probability", float("nan")) words.append( Word( text=normalized, raw=w.word, start=w.start, end=w.end, probability=float(prob) if prob is not None else float("nan"), ) ) return words, "".join(full_text).strip() def load_model( size: str = "large-v3", device: str = "cpu", compute_type: str = "int8", download_root: str = "models/whisper", ) -> WhisperModel: return WhisperModel( size, device=device, compute_type=compute_type, download_root=download_root ) def tokenize_transcript(text: str) -> list[str]: """Split a free-form transcript into normalized word tokens (lowercase, no punct).""" return [t for t in (_normalize(tok) for tok in text.split()) if t] def _char_dist(a: str, b: str) -> float: """Normalized character-level edit distance in [0, 1]. 0 = identical.""" if not a and not b: return 0.0 return 1.0 - difflib.SequenceMatcher(a=a, b=b, autojunk=False).ratio() def align_to_groundtruth( whisper_words: list[Word], gt_tokens: list[str], gap_penalty: float = 0.7 ) -> list[Word]: """Replace Whisper word identities with ground-truth tokens, keeping Whisper timing. Uses Needleman–Wunsch alignment with character-level edit distance as substitution cost. Behavior: - Aligned (W_i, GT_j): emit Word(text=GT_j, start/end from W_i). - GT_j unmatched: try to interpolate timing from neighboring aligned tokens; if impossible (start/end of utterance with no anchor), drop with a stderr warning. - W_i unmatched: drop (Whisper inserted a phantom token). """ if not whisper_words: return [] if not gt_tokens: return list(whisper_words) n, m = len(whisper_words), len(gt_tokens) dp = [[0.0] * (m + 1) for _ in range(n + 1)] bt = [[""] * (m + 1) for _ in range(n + 1)] for i in range(1, n + 1): dp[i][0] = i * gap_penalty bt[i][0] = "D" for j in range(1, m + 1): dp[0][j] = j * gap_penalty bt[0][j] = "I" for i in range(1, n + 1): for j in range(1, m + 1): sub = dp[i - 1][j - 1] + _char_dist(whisper_words[i - 1].text, gt_tokens[j - 1]) dele = dp[i - 1][j] + gap_penalty ins = dp[i][j - 1] + gap_penalty best = min(sub, dele, ins) dp[i][j] = best bt[i][j] = "M" if best == sub else ("D" if best == dele else "I") # backtrack pairs: list[tuple[int | None, int | None]] = [] i, j = n, m while i > 0 or j > 0: op = bt[i][j] if op == "M": pairs.append((i - 1, j - 1)) i -= 1 j -= 1 elif op == "D": pairs.append((i - 1, None)) i -= 1 else: pairs.append((None, j - 1)) j -= 1 pairs.reverse() # First pass: emit matched + skip Whisper-only; collect GT-only with neighbor anchors. out: list[Word | None] = [] pending_gt: list[int] = [] # GT indices waiting for next anchor last_end: float | None = None for w_idx, g_idx in pairs: if w_idx is not None and g_idx is not None: ww = whisper_words[w_idx] # if there are pending GT tokens, distribute them between last_end and ww.start if pending_gt: if last_end is not None: span_start, span_end = last_end, ww.start n_pend = len(pending_gt) for k, gi in enumerate(pending_gt): ts = span_start + (span_end - span_start) * k / n_pend te = span_start + (span_end - span_start) * (k + 1) / n_pend out.append(Word(text=gt_tokens[gi], raw=gt_tokens[gi], start=ts, end=te)) else: # at start of utterance with no anchor — assign 0..ww.start span n_pend = len(pending_gt) for k, gi in enumerate(pending_gt): ts = ww.start * k / n_pend te = ww.start * (k + 1) / n_pend out.append(Word(text=gt_tokens[gi], raw=gt_tokens[gi], start=ts, end=te)) pending_gt = [] out.append(Word(text=gt_tokens[g_idx], raw=gt_tokens[g_idx], start=ww.start, end=ww.end)) last_end = ww.end elif w_idx is not None and g_idx is None: # Whisper-only — skip, but advance time anchor last_end = whisper_words[w_idx].end else: # GT-only pending_gt.append(g_idx) # type: ignore[arg-type] # Trailing pending GT tokens (no right-anchor): assign small constant duration if pending_gt: if last_end is not None: span_start = last_end span_end = last_end + 0.3 * len(pending_gt) # rough fallback n_pend = len(pending_gt) for k, gi in enumerate(pending_gt): ts = span_start + (span_end - span_start) * k / n_pend te = span_start + (span_end - span_start) * (k + 1) / n_pend out.append(Word(text=gt_tokens[gi], raw=gt_tokens[gi], start=ts, end=te)) return [w for w in out if w is not None]