test_real_audio.py

tests/integration/test_real_audio.py

1. """Integration tier: real models on real speech. Not run in CI (model downloads).
3. TP_INTEGRATION=1 pytest tests/integration # local engine + alignment
4. TP_INTEGRATION=1 HF_TOKEN=hf_... pytest tests/integration # adds real pyannote
6. The clip is four LibriSpeech utterances from two readers (tests/fixtures), so
7. the words and the speaker boundaries are known.
8. """
10. import json
11. import os
12. from pathlib import Path
14. import pytest
16. from transcription_pipeline import steps
17. from transcription_pipeline.pipeline import Options, process
18. from transcription_pipeline.secrets import get_secret
20. pytestmark = pytest.mark.skipif(os.environ.get("TP_INTEGRATION") != "1",
21. reason="set TP_INTEGRATION=1 to run the model tier")
23. FIXTURES = Path(__file__).resolve().parents[1] / "fixtures"
24. CLIP = FIXTURES / "two-speakers.flac"
25. TRUTH = json.loads((FIXTURES / "two-speakers.json").read_text())["turns"]
28. def norm(text):
29. # The model hears LibriSpeech's "ANGOR" as "anger"; either is right.
30. return [w.strip(".,?!").replace("angor", "anger") for w in text.lower().split()]
33. def check(res, expect_speakers):
34. data = json.loads(res.outputs["json"].read_text())
35. sents = data["sentences"]
36. words = [w for s in data["segments"] for w in s.get("words", []) if "start" in w]
37. got = [w for s in sents for w in norm(s["text"])]
38. assert [w for t in TRUTH for w in norm(t["text"]) if w not in got] == []
39. starts = [w["start"] for w in words]
40. assert starts == sorted(starts) and words[-1]["end"] <= res.duration + 0.05
41. if not expect_speakers:
42. return sents
43. assert [s["speaker"] for s in sents] == ["Speaker 1", "Speaker 2", "Speaker 1", "Speaker 2"]
44. for s, t in zip(sents, TRUTH):
45. assert norm(s["text"])[0] == norm(t["text"])[0]
46. assert t["start"] - 0.3 <= s["start"] <= t["end"]
47. return sents
50. def test_local_engine_and_alignment(tmp_path):
51. opts = Options(language="en", diarize=False)
52. res = process(CLIP, tmp_path, opts, steps.transcriber(opts), steps.aligner(opts), None)
53. check(res, expect_speakers=False)
56. def test_speaker_assignment_on_real_words(tmp_path):
57. """Real transcription, alignment and whisperx.assign_word_speakers; only the
58. diarizer is replaced by the known turns."""
59. import pandas as pd
60. import whisperx
61. turns = pd.DataFrame([{"start": t["start"], "end": t["end"],
62. "speaker": f"SPEAKER_0{'AB'.index(t['speaker'])}"} for t in TRUTH])
64. def known(result, wav, o):
65. return whisperx.assign_word_speakers(turns, result)
66. opts = Options(language="en")
67. res = process(CLIP, tmp_path, opts, steps.transcriber(opts), steps.aligner(opts), known)
68. check(res, expect_speakers=True)
71. @pytest.mark.skipif(not get_secret("HF_TOKEN"), reason="needs a Hugging Face token for pyannote")
72. def test_pyannote_diarization(tmp_path):
73. opts = Options(language="en", min_speakers=2, max_speakers=2)
74. res = process(CLIP, tmp_path, opts, steps.transcriber(opts), steps.aligner(opts),
75. steps.diarizer(opts))
76. check(res, expect_speakers=True)