diff --git a/diarize.py b/diarize.py new file mode 100644 index 0000000..4240ef6 --- /dev/null +++ b/diarize.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +"""Optional speaker-diarization hook for codec-carver. + +This module answers "who spoke when?" for an audio file and lets callers +attribute transcript segments to speakers. It follows the repository's +optional-AI pattern: + +- **Dependency-free at import.** Importing this module never pulls in heavy + ML libraries; only the standard library is touched. +- **Injectable backend.** ``diarize_file`` accepts any callable that maps an + audio path to speaker turns, so tests and alternative engines never need + the real model. +- **Lazy heavy import.** The default backend imports ``pyannote.audio`` only + when actually invoked. +- **Graceful degradation.** When the optional dependency is missing, a + :class:`DiarizationUnavailableError` with an actionable install message is + raised instead of a bare ``ImportError`` at import time. + +The pure functions :func:`merge_with_transcript` and :func:`to_text` work +without any model at all, so speaker attribution logic stays fully testable +offline. + +Honest note on the optional dependency: the default backend relies on +``pyannote.audio``, a large package that downloads pretrained models on +first use and requires a Hugging Face access token for the gated +``pyannote/speaker-diarization-3.1`` pipeline. Nothing in codec-carver +requires it unless you call :func:`diarize_file` without a custom backend. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable, Iterable, Sequence + + +DEFAULT_PIPELINE_NAME = "pyannote/speaker-diarization-3.1" + +FALLBACK_SPEAKER = "SPEAKER_1" + +_INSTALL_HINT = ( + "Speaker diarization requires the optional 'pyannote.audio' package, " + "which is not installed. Install it with:\n" + " pip install pyannote.audio\n" + "You will also need a Hugging Face access token with permission for " + f"the gated '{DEFAULT_PIPELINE_NAME}' pipeline " + "(set HF_TOKEN or pass hf_token). Alternatively, pass your own " + "backend callable to diarize_file(..., backend=...)." +) + + +class DiarizationUnavailableError(RuntimeError): + """Raised when diarization is requested but no backend can run. + + This typically means the optional ``pyannote.audio`` dependency is not + installed (or its pretrained pipeline could not be loaded). The error + message always includes actionable installation instructions. + """ + + +@dataclass(frozen=True) +class SpeakerTurn: + """A contiguous interval during which a single speaker is talking. + + Attributes: + start: Interval start in seconds from the beginning of the audio. + end: Interval end in seconds from the beginning of the audio. + speaker: Backend-assigned speaker label (e.g. ``"SPEAKER_00"``). + """ + + start: float + end: float + speaker: str + + +@dataclass(frozen=True) +class DiarizationResult: + """The outcome of diarizing one audio file. + + Attributes: + turns: Chronologically ordered speaker turns. + speaker_count: Number of distinct speaker labels across all turns. + """ + + turns: tuple[SpeakerTurn, ...] = field(default_factory=tuple) + speaker_count: int = 0 + + +@dataclass(frozen=True) +class AttributedSegment: + """A transcript segment annotated with its most likely speaker. + + Attributes: + start: Segment start in seconds. + end: Segment end in seconds. + text: Transcribed text of the segment. + speaker: Speaker label with the greatest time overlap, or + :data:`FALLBACK_SPEAKER` when no turn overlaps the segment. + """ + + start: float + end: float + text: str + speaker: str + + +def _default_backend(audio_path: str) -> list[SpeakerTurn]: + """Diarize ``audio_path`` with ``pyannote.audio`` (imported lazily). + + The import happens inside this function so that merely importing + :mod:`diarize` never requires the heavy optional dependency. + + Args: + audio_path: Path to the audio file to diarize. + + Returns: + Speaker turns produced by the pretrained pipeline. + + Raises: + DiarizationUnavailableError: If ``pyannote.audio`` is not installed + or the pretrained pipeline cannot be loaded. + """ + try: + from pyannote.audio import Pipeline + except ImportError as exc: + raise DiarizationUnavailableError(_INSTALL_HINT) from exc + + try: + pipeline = Pipeline.from_pretrained(DEFAULT_PIPELINE_NAME) + except Exception as exc: # pragma: no cover - depends on remote model + raise DiarizationUnavailableError( + f"Could not load the '{DEFAULT_PIPELINE_NAME}' pipeline: {exc}\n" + + _INSTALL_HINT + ) from exc + + annotation = pipeline(audio_path) + return [ + SpeakerTurn(start=segment.start, end=segment.end, speaker=str(label)) + for segment, _, label in annotation.itertracks(yield_label=True) + ] + + +def diarize_file( + audio_path: str, + *, + backend: Callable[[str], Iterable[SpeakerTurn]] | None = None, +) -> DiarizationResult: + """Identify who speaks when in ``audio_path``. + + Args: + audio_path: Path to the audio file to diarize. + backend: Optional callable mapping an audio path to an iterable of + :class:`SpeakerTurn`. When omitted, the default backend lazily + imports ``pyannote.audio``; if that optional dependency is + missing, :class:`DiarizationUnavailableError` is raised with + install instructions. + + Returns: + A :class:`DiarizationResult` whose turns are sorted by start time + (ties broken by end time) and whose ``speaker_count`` is the number + of distinct speaker labels. + + Raises: + DiarizationUnavailableError: If the default backend is selected and + ``pyannote.audio`` is unavailable. + """ + runner = backend if backend is not None else _default_backend + turns = tuple(sorted(runner(audio_path), key=lambda t: (t.start, t.end))) + speakers = {turn.speaker for turn in turns} + return DiarizationResult(turns=turns, speaker_count=len(speakers)) + + +def _overlap(a_start: float, a_end: float, b_start: float, b_end: float) -> float: + """Return the length in seconds of the overlap between two intervals. + + Args: + a_start: Start of the first interval. + a_end: End of the first interval. + b_start: Start of the second interval. + b_end: End of the second interval. + + Returns: + The positive overlap duration, or ``0.0`` when the intervals are + disjoint (or merely touch at a boundary). + """ + return max(0.0, min(a_end, b_end) - max(a_start, b_start)) + + +def merge_with_transcript( + turns: Iterable[SpeakerTurn], + segments: Iterable[object], +) -> list[AttributedSegment]: + """Assign a speaker to each transcript segment by maximum time overlap. + + This is a pure function: no model, file access, or network is involved, + so it is fully testable offline. + + Args: + turns: Speaker turns, e.g. from :func:`diarize_file`. + segments: Duck-typed transcript segments exposing ``.start``, + ``.end`` (seconds) and ``.text`` attributes — for example the + segment objects produced by a transcription tool. + + Returns: + One :class:`AttributedSegment` per input segment, in input order. + Each segment gets the speaker whose turns overlap it the longest in + total; when several speakers tie, the one whose overlapping turn + appears earliest wins. Segments that overlap no turn at all fall + back to :data:`FALLBACK_SPEAKER`. + """ + turn_list = list(turns) + merged: list[AttributedSegment] = [] + for segment in segments: + totals: dict[str, float] = {} + for turn in turn_list: + shared = _overlap(segment.start, segment.end, turn.start, turn.end) + if shared > 0.0: + totals[turn.speaker] = totals.get(turn.speaker, 0.0) + shared + if totals: + speaker = max(totals, key=lambda name: totals[name]) + else: + speaker = FALLBACK_SPEAKER + merged.append( + AttributedSegment( + start=segment.start, + end=segment.end, + text=segment.text, + speaker=speaker, + ) + ) + return merged + + +def to_text(merged: Sequence[AttributedSegment]) -> str: + """Render speaker-attributed segments as human-readable lines. + + Args: + merged: Attributed segments, e.g. from :func:`merge_with_transcript`. + + Returns: + One ``"[speaker] text"`` line per segment, joined by newlines. An + empty input yields an empty string. + """ + return "\n".join(f"[{item.speaker}] {item.text}" for item in merged) diff --git a/tests/test_diarize.py b/tests/test_diarize.py new file mode 100644 index 0000000..ecccb3f --- /dev/null +++ b/tests/test_diarize.py @@ -0,0 +1,217 @@ +"""Tests for the optional speaker-diarization hook (diarize.py). + +All tests run offline: backends are injected fakes, and the unavailable +path is exercised by blocking the optional import via sys.modules. No +model download or network access ever happens here. +""" + +import sys +import unittest +from dataclasses import dataclass +from pathlib import Path +from unittest.mock import patch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import diarize +from diarize import ( + AttributedSegment, + DiarizationResult, + DiarizationUnavailableError, + SpeakerTurn, + diarize_file, + merge_with_transcript, + to_text, +) + + +@dataclass +class FakeSegment: + """Duck-typed transcript segment with .start/.end/.text attributes.""" + + start: float + end: float + text: str + + +class TestImportIsLight(unittest.TestCase): + """Importing diarize must not drag in the optional heavy dependency.""" + + def test_pyannote_not_imported_by_module_import(self): + self.assertIn("diarize", sys.modules) + self.assertNotIn("pyannote.audio", sys.modules) + + +class TestDiarizeFileWithInjectedBackend(unittest.TestCase): + """diarize_file with a fake backend: no model, no network.""" + + def test_returns_result_with_turns_and_speaker_count(self): + def backend(audio_path): + self.assertEqual(audio_path, "meeting.wav") + return [ + SpeakerTurn(0.0, 2.0, "SPEAKER_00"), + SpeakerTurn(2.0, 4.0, "SPEAKER_01"), + SpeakerTurn(4.0, 6.0, "SPEAKER_00"), + ] + + result = diarize_file("meeting.wav", backend=backend) + self.assertIsInstance(result, DiarizationResult) + self.assertEqual(len(result.turns), 3) + self.assertEqual(result.speaker_count, 2) + + def test_turns_are_sorted_by_start_time(self): + def backend(_): + return [ + SpeakerTurn(5.0, 6.0, "B"), + SpeakerTurn(0.0, 1.0, "A"), + SpeakerTurn(2.0, 3.0, "C"), + ] + + result = diarize_file("x.wav", backend=backend) + self.assertEqual([t.start for t in result.turns], [0.0, 2.0, 5.0]) + + def test_empty_backend_output(self): + result = diarize_file("silent.wav", backend=lambda _: []) + self.assertEqual(result.turns, ()) + self.assertEqual(result.speaker_count, 0) + + def test_backend_may_return_any_iterable(self): + result = diarize_file( + "gen.wav", + backend=lambda _: iter([SpeakerTurn(0.0, 1.0, "S")]), + ) + self.assertEqual(result.speaker_count, 1) + + +class TestDefaultBackendUnavailable(unittest.TestCase): + """Without pyannote.audio installed, the error must be clear and actionable.""" + + def _call_with_import_blocked(self): + blocked = {"pyannote": None, "pyannote.audio": None} + with patch.dict(sys.modules, blocked): + diarize_file("audio.wav") + + def test_raises_diarization_unavailable_error(self): + with self.assertRaises(DiarizationUnavailableError): + self._call_with_import_blocked() + + def test_error_message_is_actionable(self): + with self.assertRaises(DiarizationUnavailableError) as ctx: + self._call_with_import_blocked() + message = str(ctx.exception) + self.assertIn("pip install pyannote.audio", message) + self.assertIn("backend", message) + + def test_error_chains_the_original_import_error(self): + with self.assertRaises(DiarizationUnavailableError) as ctx: + self._call_with_import_blocked() + self.assertIsInstance(ctx.exception.__cause__, ImportError) + + +class TestMergeWithTranscript(unittest.TestCase): + """Pure overlap-based speaker assignment.""" + + def test_exact_overlap_assignment(self): + turns = [ + SpeakerTurn(0.0, 5.0, "SPEAKER_00"), + SpeakerTurn(5.0, 10.0, "SPEAKER_01"), + ] + segments = [ + FakeSegment(0.0, 5.0, "hello"), + FakeSegment(5.0, 10.0, "world"), + ] + merged = merge_with_transcript(turns, segments) + self.assertEqual([m.speaker for m in merged], ["SPEAKER_00", "SPEAKER_01"]) + self.assertEqual([m.text for m in merged], ["hello", "world"]) + + def test_partial_overlap_picks_max_overlap_speaker(self): + turns = [ + SpeakerTurn(0.0, 4.0, "SPEAKER_00"), + SpeakerTurn(4.0, 10.0, "SPEAKER_01"), + ] + # Segment spans 3..9: 1s with SPEAKER_00, 5s with SPEAKER_01. + merged = merge_with_transcript(turns, [FakeSegment(3.0, 9.0, "mixed")]) + self.assertEqual(merged[0].speaker, "SPEAKER_01") + + def test_no_overlap_falls_back_to_speaker_1(self): + turns = [SpeakerTurn(0.0, 1.0, "SPEAKER_00")] + merged = merge_with_transcript(turns, [FakeSegment(50.0, 55.0, "late")]) + self.assertEqual(merged[0].speaker, "SPEAKER_1") + + def test_empty_turns_falls_back_to_speaker_1(self): + merged = merge_with_transcript([], [FakeSegment(0.0, 1.0, "solo")]) + self.assertEqual(merged[0].speaker, "SPEAKER_1") + + def test_touching_boundary_is_not_overlap(self): + turns = [SpeakerTurn(0.0, 2.0, "SPEAKER_00")] + merged = merge_with_transcript(turns, [FakeSegment(2.0, 4.0, "after")]) + self.assertEqual(merged[0].speaker, "SPEAKER_1") + + def test_multi_speaker_accumulates_split_turns(self): + # SPEAKER_00 overlaps in two short turns (1s + 2s = 3s total), + # SPEAKER_01 overlaps once for 2.5s; accumulation must win for 00. + turns = [ + SpeakerTurn(0.0, 1.0, "SPEAKER_00"), + SpeakerTurn(1.0, 3.5, "SPEAKER_01"), + SpeakerTurn(3.5, 5.5, "SPEAKER_00"), + ] + merged = merge_with_transcript(turns, [FakeSegment(0.0, 5.5, "debate")]) + self.assertEqual(merged[0].speaker, "SPEAKER_00") + + def test_multiple_segments_multi_speaker(self): + turns = [ + SpeakerTurn(0.0, 3.0, "A"), + SpeakerTurn(3.0, 6.0, "B"), + SpeakerTurn(6.0, 9.0, "C"), + ] + segments = [ + FakeSegment(0.5, 2.5, "one"), + FakeSegment(3.5, 5.5, "two"), + FakeSegment(6.5, 8.5, "three"), + ] + merged = merge_with_transcript(turns, segments) + self.assertEqual([m.speaker for m in merged], ["A", "B", "C"]) + + def test_preserves_segment_order_and_fields(self): + turns = [SpeakerTurn(0.0, 10.0, "S")] + segments = [FakeSegment(1.0, 2.0, "first"), FakeSegment(3.0, 4.0, "second")] + merged = merge_with_transcript(turns, segments) + self.assertEqual(merged[0].start, 1.0) + self.assertEqual(merged[0].end, 2.0) + self.assertEqual([m.text for m in merged], ["first", "second"]) + + def test_empty_segments_returns_empty_list(self): + self.assertEqual(merge_with_transcript([SpeakerTurn(0, 1, "S")], []), []) + + +class TestToText(unittest.TestCase): + """Rendering attributed segments as '[speaker] text' lines.""" + + def test_formats_speaker_prefixed_lines(self): + merged = [ + AttributedSegment(0.0, 1.0, "hello there", "SPEAKER_00"), + AttributedSegment(1.0, 2.0, "hi back", "SPEAKER_01"), + ] + self.assertEqual( + to_text(merged), + "[SPEAKER_00] hello there\n[SPEAKER_01] hi back", + ) + + def test_empty_input_yields_empty_string(self): + self.assertEqual(to_text([]), "") + + def test_round_trip_with_merge(self): + turns = [SpeakerTurn(0.0, 2.0, "SPEAKER_00")] + merged = merge_with_transcript(turns, [FakeSegment(0.0, 2.0, "solo line")]) + self.assertEqual(to_text(merged), "[SPEAKER_00] solo line") + + +class TestModuleConstants(unittest.TestCase): + """Public constants exposed for callers.""" + + def test_fallback_speaker_constant(self): + self.assertEqual(diarize.FALLBACK_SPEAKER, "SPEAKER_1") + + +if __name__ == "__main__": + unittest.main()