added pyannote

app/services/asr.py

@@ -1,51 +1,89 @@
 from typing import List, Dict
 from faster_whisper import WhisperModel
 import os
+import torch
+from pyannote.audio import Pipeline
 
 # Load environment variables
 WHISPER_MODEL_SIZE = os.getenv("WHISPER_MODEL_SIZE", "base")
-WHISPER_DEVICE = os.getenv("WHISPER_DEVICE", "cpu")
-WHISPER_COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "int8")
+HF_TOKEN = os.getenv("HF_TOKEN")  # For pyannote
 
+# Initialize Whisper model
 try:
-    _whisper = WhisperModel(
-        WHISPER_MODEL_SIZE,
-        device="cpu",  # Force CPU to avoid CUDA issues
-        compute_type="int8"
-    )
+    _whisper = WhisperModel(WHISPER_MODEL_SIZE, device="cpu", compute_type="int8")
 except Exception as e:
     print(f"Error loading Whisper model: {e}")
-    # Fallback to base model with explicit CPU settings
-    _whisper = WhisperModel(
-        "base",
-        device="cpu",
-        compute_type="int8"
-    )
+    _whisper = WhisperModel("base", device="cpu", compute_type="int8")
+
+# Initialize pyannote diarization pipeline
+if HF_TOKEN:
+    try:
+        diarization_pipeline = Pipeline.from_pretrained(
+            "pyannote/speaker-diarization-3.1", use_auth_token=HF_TOKEN
+        )
+        # Move pipeline to CPU if no GPU is available
+        if not torch.cuda.is_available():
+            diarization_pipeline = diarization_pipeline.to(torch.device("cpu"))
+    except Exception as e:
+        print(f"Error loading pyannote pipeline: {e}")
+        diarization_pipeline = None
+else:
+    print("HUGGING_FACE_TOKEN not set, skipping diarization.")
+    diarization_pipeline = None
+
 
 def transcribe(audio_path: str) -> List[Dict]:
     """
-    Transcribe an audio file using faster-whisper and return our standard
-    list of segments: [{start, end, speaker, text}].
+    Transcribe an audio file using faster-whisper and combine with
+    pyannote.audio for speaker diarization.
     """
-    # beam_size=1 is fastest; raise for a bit more accuracy.
-    segments, info = _whisper.transcribe(
-        audio_path,
-        language="en",         
-        vad_filter=True,      
-        beam_size=1
+    # 1. Transcribe with Whisper
+    segments, _ = _whisper.transcribe(
+        audio_path, language="en", vad_filter=True, beam_size=1
     )
 
-    out = []
+    whisper_segments = []
     for seg in segments:
-        out.append({
-            "start": float(seg.start),
-            "end": float(seg.end),
-            "speaker": "A",    # no diarization here; we can add later
-            "text": seg.text.strip()
-        })
-
-    # If there were no segments (edge case), return a single empty segment
-    if not out:
-        out = [{"start": 0.0, "end": 0.0, "speaker": "A", "text": ""}]
-
-    return out
+        whisper_segments.append(
+            {"start": float(seg.start), "end": float(seg.end), "text": seg.text.strip()}
+        )
+
+    if not diarization_pipeline:
+        # If diarization is not available, return with a single speaker
+        for seg in whisper_segments:
+            seg["speaker"] = "A"
+        return whisper_segments
+
+    # 2. Perform Diarization
+    try:
+        diarization = diarization_pipeline(audio_path)
+    except Exception as e:
+        print(f"Error during diarization: {e}")
+        for seg in whisper_segments:
+            seg["speaker"] = "A"
+        return whisper_segments
+
+    # 3. Assign Speaker to Segments
+    out_segments = []
+    for seg in whisper_segments:
+        # Find the speaker for the segment's midpoint
+        midpoint = seg["start"] + (seg["end"] - seg["start"]) / 2
+        speaker = "UNKNOWN"
+        for turn, _, speaker_label in diarization.itertracks(yield_label=True):
+            if turn.start <= midpoint <= turn.end:
+                speaker = speaker_label
+                break
+
+        out_segments.append(
+            {
+                "start": seg["start"],
+                "end": seg["end"],
+                "speaker": speaker,
+                "text": seg["text"],
+            }
+        )
+
+    if not out_segments:
+        return [{"start": 0.0, "end": 0.0, "speaker": "A", "text": ""}]
+
+    return out_segments

requirements.txt

@@ -62,3 +62,6 @@ uvicorn==0.35.0
 uvloop==0.21.0
 watchfiles==1.1.0
 websockets==15.0.1
+
+# Added for Speaker Diarization
+pyannote.audio>=3.3.2