align.py

transcription_pipeline/align.py

1. """Word timestamps: WhisperX's wav2vec2 forced alignment.
3. Runs on either engine's segments. Neither Whisper model returns reliable word
4. timing on its own (DeepInfra returns none), and speaker assignment needs it.
5. When no alignment model exists for the language, the result passes through
6. untouched and shape.py interpolates sentence times across each segment.
7. """
9. from __future__ import annotations
11. from functools import lru_cache
12. from pathlib import Path
14. from .audio import load_wav16k
17. @lru_cache(maxsize=4)
18. def _align_model(language: str, device: str):
19. import whisperx
20. return whisperx.load_align_model(language_code=language, device=device)
23. def align(result: dict, wav: Path, device: str = "cpu", log=print) -> dict:
24. import whisperx
25. language = result.get("language") or "en"
26. try:
27. model, metadata = _align_model(language, device)
28. except ValueError as e: # no default alignment model for this language
29. log(f" align: skipped ({e}); timestamps will be interpolated.")
30. return result
31. audio = load_wav16k(wav)
32. aligned = whisperx.align(result["segments"], model, metadata, audio, device,
33. return_char_alignments=False)
34. return {**result, "segments": aligned["segments"], "aligned": True}