Update pipeline.py

pipeline.py

@@ -1,7 +1,7 @@
 """
-Speech Annotation Pipeline
-- Speaker Diarization via pyannote.audio
-- Transcription via openai-whisper
+Speech Annotation Pipeline using Deepgram
+- Transcription + Speaker Diarization in one API call
+- Much faster and more accurate than pyannote on mixed audio
 """
 
 import os
@@ -9,25 +9,15 @@ import json
 import csv
 import time
 import logging
+import httpx
 from pathlib import Path
 from dataclasses import dataclass
 from typing import List, Optional
 
-import torch
-import whisper
-from huggingface_hub import login
-from pyannote.audio import Pipeline as PyannotePipeline
-
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-HF_TOKEN = os.environ.get("HF_TOKEN", "")
-WHISPER_MODEL = os.environ.get("WHISPER_MODEL", "base")
-NUM_SPEAKERS = int(os.environ.get("NUM_SPEAKERS", "0"))  # 0 = auto-detect
-
-if HF_TOKEN:
-    login(token=HF_TOKEN)
+DEEPGRAM_API_KEY = os.environ.get("DEEPGRAM_API_KEY", "")
 
 
 @dataclass
@@ -51,102 +41,110 @@ class AnnotationSegment:
         return f"{self.speaker} {self.start_fmt} {self.text}"
 
 
-def _build_speaker_map(raw_labels):
-    letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
-    seen, mapping = [], {}
-    for lbl in raw_labels:
-        if lbl not in seen:
-            seen.append(lbl)
-    for i, lbl in enumerate(seen):
-        mapping[lbl] = letters[i] if i < len(letters) else f"SPK{i}"
-    return mapping
+SPEAKER_LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+
+def _speaker_label(idx: int) -> str:
+    return SPEAKER_LETTERS[idx] if idx < len(SPEAKER_LETTERS) else f"SPK{idx}"
 
 
 class SpeechAnnotationPipeline:
-    def __init__(self):
-        self._whisper_model = None
-        self._diarization_pipeline = None
-
-    def _load_whisper(self):
-        if self._whisper_model is None:
-            logger.info(f"Loading Whisper model '{WHISPER_MODEL}' on {DEVICE}...")
-            self._whisper_model = whisper.load_model(WHISPER_MODEL, device=DEVICE)
-        return self._whisper_model
-
-    def _load_diarization(self):
-        if self._diarization_pipeline is None:
-            if not HF_TOKEN:
-                raise EnvironmentError(
-                    "HF_TOKEN not set. Add it in Space Settings > Variables and secrets."
-                )
-            logger.info("Loading pyannote diarization pipeline...")
-            self._diarization_pipeline = PyannotePipeline.from_pretrained(
-                "pyannote/speaker-diarization-3.1",
-                use_auth_token=HF_TOKEN,
+
+    def process(self, audio_path: str, num_speakers: int = 0, **kwargs) -> List[AnnotationSegment]:
+        if not DEEPGRAM_API_KEY:
+            raise EnvironmentError(
+                "DEEPGRAM_API_KEY not set. Add it in Space Settings > Variables and secrets."
             )
-            if DEVICE == "cuda":
-                self._diarization_pipeline = self._diarization_pipeline.to(torch.device("cuda"))
-        return self._diarization_pipeline
-
-    def _transcribe(self, audio_path):
-        logger.info("Step 1/2 - Transcribing with Whisper...")
-        model = self._load_whisper()
-        result = model.transcribe(audio_path, verbose=False)
-        return result["segments"]
-
-    def _diarize(self, audio_path, num_speakers=0, min_speakers=None, max_speakers=None):
-        logger.info("Step 2/2 - Diarizing speakers...")
-        pipeline = self._load_diarization()
-
-        # Build diarization params
-        params = {}
+
+        logger.info(f"Sending audio to Deepgram: {audio_path}")
+        t0 = time.time()
+
+        # Read audio file
+        audio_bytes = Path(audio_path).read_bytes()
+        ext = Path(audio_path).suffix.lower().lstrip(".")
+        mime = {
+            "wav": "audio/wav",
+            "mp3": "audio/mpeg",
+            "mp4": "audio/mp4",
+            "m4a": "audio/mp4",
+            "flac": "audio/flac",
+            "ogg": "audio/ogg",
+        }.get(ext, "audio/wav")
+
+        # Build Deepgram request params
+        params = {
+            "model": "nova-2",
+            "diarize": "true",
+            "punctuate": "true",
+            "utterances": "true",
+            "smart_format": "true",
+        }
         if num_speakers and num_speakers > 0:
-            params["num_speakers"] = num_speakers
-            logger.info(f"Using fixed num_speakers={num_speakers}")
-        else:
-            if min_speakers:
-                params["min_speakers"] = min_speakers
-            if max_speakers:
-                params["max_speakers"] = max_speakers
-
-        diarization = pipeline(audio_path, **params)
-        segments = []
-        for turn, _, speaker in diarization.itertracks(yield_label=True):
-            segments.append({"start": turn.start, "end": turn.end, "speaker": speaker})
-        return segments
+            params["diarize_version"] = "latest"
+
+        headers = {
+            "Authorization": f"Token {DEEPGRAM_API_KEY}",
+            "Content-Type": mime,
+        }
+
+        # Call Deepgram API
+        response = httpx.post(
+            "https://api.deepgram.com/v1/listen",
+            params=params,
+            headers=headers,
+            content=audio_bytes,
+            timeout=300,
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        # Parse utterances (already speaker-separated)
+        utterances = data.get("results", {}).get("utterances", [])
+        if not utterances:
+            # Fallback: parse words
+            words = data.get("results", {}).get("channels", [{}])[0]\
+                        .get("alternatives", [{}])[0].get("words", [])
+            utterances = self._words_to_utterances(words)
 
-    def _assign_speakers(self, transcript_segments, diarization_segments):
-        raw_labels = [d["speaker"] for d in diarization_segments]
-        speaker_map = _build_speaker_map(raw_labels)
-        results = []
-        for seg in transcript_segments:
-            t_start, t_end = seg["start"], seg["end"]
-            text = seg.get("text", "").strip()
-            if not text:
-                continue
-            best_speaker, best_overlap = "A", 0.0
-            for d in diarization_segments:
-                overlap = max(0, min(t_end, d["end"]) - max(t_start, d["start"]))
-                if overlap > best_overlap:
-                    best_overlap = overlap
-                    best_speaker = d["speaker"]
-            results.append(AnnotationSegment(
-                speaker=speaker_map.get(best_speaker, best_speaker),
-                start=t_start, end=t_end, text=text,
+        segments = []
+        for utt in utterances:
+            speaker_idx = utt.get("speaker", 0)
+            segments.append(AnnotationSegment(
+                speaker=_speaker_label(speaker_idx),
+                start=utt.get("start", 0),
+                end=utt.get("end", 0),
+                text=utt.get("transcript", "").strip(),
             ))
-        return results
 
-    def process(self, audio_path, num_speakers=0, min_speakers=None, max_speakers=None):
-        t0 = time.time()
-        transcript = self._transcribe(audio_path)
-        diarization = self._diarize(audio_path, num_speakers, min_speakers, max_speakers)
-        annotations = self._assign_speakers(transcript, diarization)
-        logger.info(f"Done in {time.time()-t0:.1f}s — {len(annotations)} segments")
-        return annotations
+        logger.info(f"Done in {time.time()-t0:.1f}s — {len(segments)} segments")
+        return segments
+
+    def _words_to_utterances(self, words):
+        """Group words into utterances by speaker."""
+        if not words:
+            return []
+        utterances = []
+        current = {"speaker": words[0].get("speaker", 0),
+                   "start": words[0].get("start", 0),
+                   "end": words[0].get("end", 0),
+                   "transcript": words[0].get("punctuated_word", words[0].get("word", ""))}
+        for word in words[1:]:
+            if word.get("speaker") == current["speaker"]:
+                current["end"] = word.get("end", current["end"])
+                current["transcript"] += " " + word.get("punctuated_word", word.get("word", ""))
+            else:
+                utterances.append(current)
+                current = {
+                    "speaker": word.get("speaker", 0),
+                    "start": word.get("start", 0),
+                    "end": word.get("end", 0),
+                    "transcript": word.get("punctuated_word", word.get("word", "")),
+                }
+        utterances.append(current)
+        return utterances
 
 
 def to_json(segments, path):
-    data = [{"speaker": s.speaker, "start": round(s.start,3), "end": round(s.end,3),
+    data = [{"speaker": s.speaker, "start": round(s.start, 3), "end": round(s.end, 3),
               "start_fmt": s.start_fmt, "end_fmt": s.end_fmt, "text": s.text} for s in segments]
     with open(path, "w", encoding="utf-8") as f:
         json.dump(data, f, indent=2, ensure_ascii=False)
@@ -161,9 +159,9 @@ def to_csv(segments, path):
 def to_abab_text(segments):
     return "\n".join(s.to_abab() for s in segments)
 
-_pipeline_instance = None
+_pipeline_instance: Optional[SpeechAnnotationPipeline] = None
 
-def get_pipeline():
+def get_pipeline() -> SpeechAnnotationPipeline:
     global _pipeline_instance
     if _pipeline_instance is None:
         _pipeline_instance = SpeechAnnotationPipeline()