diarize.py

transcription_pipeline/diarize.py

1. """Who spoke when: pyannote speaker diarization through WhisperX.
3. The default model is WhisperX's own default, pyannote/speaker-diarization-community-1
4. (the successor to speaker-diarization-3.1); set TP_DIARIZE_MODEL to use another.
5. It needs a Hugging Face token whose account has accepted that model's terms
6. (docs/SETUP.md). The model downloads once, then runs locally. Each aligned word
7. is given the speaker whose turn overlaps it most; shape.py turns that into
8. sentences.
9. """
11. from __future__ import annotations
13. import inspect
14. import os
15. from functools import lru_cache
16. from pathlib import Path
18. from .audio import load_wav16k
20. DEFAULT_MODEL = os.environ.get("TP_DIARIZE_MODEL") or "pyannote/speaker-diarization-community-1"
23. class DiarizeError(RuntimeError):
24. pass
27. def explain_load_error(err: Exception, model: str) -> DiarizeError | None:
28. """Turn Hugging Face's access errors into what to do about them."""
29. name = type(err).__name__
30. text = str(err)
31. if name == "GatedRepoError" or "gated repo" in text.lower() or "authorized list" in text:
32. return DiarizeError(
33. f"Your Hugging Face account has not accepted the terms for {model}. "
34. f"Open https://huggingface.co/{model} while signed in, accept them, then try again "
35. "(or run with --no-diarize).")
36. if "401" in text or "Invalid credentials" in text or "invalid user token" in text.lower():
37. return DiarizeError("Hugging Face rejected the token. Make a new read token and store it "
38. "again (docs/SETUP.md, step 2).")
39. return None
42. @lru_cache(maxsize=1)
43. def _pipeline(model: str, token: str, device: str):
44. from whisperx.diarize import DiarizationPipeline
45. params = inspect.signature(DiarizationPipeline.__init__).parameters
46. # Older WhisperX called this argument use_auth_token.
47. key = "token" if "token" in params else "use_auth_token"
48. try:
49. return DiarizationPipeline(model_name=model, device=device, **{key: token})
50. except Exception as err:
51. explained = explain_load_error(err, model)
52. if explained is None:
53. raise
54. raise explained from err
57. def diarize(result: dict, wav: Path, token: str | None, min_speakers: int | None = None,
58. max_speakers: int | None = None, device: str = "cpu",
59. model: str = DEFAULT_MODEL) -> dict:
60. if not token:
61. raise DiarizeError("No Hugging Face token: set HF_TOKEN or add it to the Keychain, or run "
62. "with --no-diarize (see docs/SETUP.md).")
63. import whisperx
64. audio = load_wav16k(wav)
65. if min_speakers and min_speakers == max_speakers:
66. speakers = {"num_speakers": min_speakers}
67. else:
68. speakers = {"min_speakers": min_speakers, "max_speakers": max_speakers}
69. turns = _pipeline(model, token, device)(audio, **speakers)
70. return {**whisperx.assign_word_speakers(turns, result), "diarized": True}