Update pipeline.py

pipeline.py

@@ -1,168 +1,68 @@
-"""
-Speech Annotation Pipeline using Deepgram
-- Transcription + Speaker Diarization in one API call
-- Much faster and more accurate than pyannote on mixed audio
-"""
-
-import os
-import json
-import csv
-import time
-import logging
-import httpx
+import tempfile
+import traceback
 from pathlib import Path
-from dataclasses import dataclass
-from typing import List, Optional
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-DEEPGRAM_API_KEY = os.environ.get("DEEPGRAM_API_KEY", "")
-
-
-@dataclass
-class AnnotationSegment:
-    speaker: str
-    start: float
-    end: float
-    text: str
-
-    @property
-    def start_fmt(self):
-        m, s = divmod(int(self.start), 60)
-        return f"{m:02d}:{s:02d}"
-
-    @property
-    def end_fmt(self):
-        m, s = divmod(int(self.end), 60)
-        return f"{m:02d}:{s:02d}"
-
-    def to_abab(self):
-        return f"{self.speaker} {self.start_fmt} {self.text}"
-
-
-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 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."
-            )
-
-        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["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)
-
-        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(),
-            ))
-
-        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),
-              "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)
-
-def to_csv(segments, path):
-    with open(path, "w", newline="", encoding="utf-8") as f:
-        writer = csv.writer(f)
-        writer.writerow(["speaker", "start", "end", "text"])
-        for s in segments:
-            writer.writerow([s.speaker, s.start_fmt, s.end_fmt, s.text])
-
-def to_abab_text(segments):
-    return "\n".join(s.to_abab() for s in segments)
-
-_pipeline_instance: Optional[SpeechAnnotationPipeline] = None
-
-def get_pipeline() -> SpeechAnnotationPipeline:
-    global _pipeline_instance
-    if _pipeline_instance is None:
-        _pipeline_instance = SpeechAnnotationPipeline()
-    return _pipeline_instance
+import gradio as gr
+from pipeline import get_pipeline, to_json, to_csv, to_abab_text
+
+UPLOAD_DIR = Path(tempfile.gettempdir()) / "speech_annotation"
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+
+_last_segments = []
+
+def process_audio(audio_path, num_speakers):
+    global _last_segments
+    if audio_path is None:
+        return "⚠️ Please upload an audio file first.", ""
+    try:
+        pipeline = get_pipeline()
+        n = int(num_speakers) if num_speakers and int(num_speakers) > 0 else 0
+        segments = pipeline.process(audio_path, num_speakers=n)
+    except Exception as e:
+        return f"❌ Error: {e}\n{traceback.format_exc()}", ""
+    if not segments:
+        return "⚠️ No speech detected.", ""
+    _last_segments = segments
+    abab = to_abab_text(segments)
+    unique = len(set(s.speaker for s in segments))
+    return f"✅ Done — {len(segments)} segments, {unique} speaker(s) detected", abab
+
+def export_json():
+    if not _last_segments:
+        return None
+    out = str(UPLOAD_DIR / "annotation.json")
+    to_json(_last_segments, out)
+    return out
+
+def export_csv():
+    if not _last_segments:
+        return None
+    out = str(UPLOAD_DIR / "annotation.csv")
+    to_csv(_last_segments, out)
+    return out
+
+with gr.Blocks(title="Speech Annotation Pipeline") as demo:
+    gr.Markdown("# 🎙️ Speech Annotation Pipeline")
+    audio_input = gr.Audio(label="Upload Audio (.wav/.mp3/.flac)", type="filepath")
+    num_speakers = gr.Slider(minimum=0, maximum=10, step=1, value=0,
+                             label="Number of speakers (0 = auto-detect)")
+    run_btn = gr.Button("▶ Run Annotation", variant="primary")
+    status_box = gr.Textbox(label="Status", value="Ready.", interactive=False)
+    abab_box = gr.Textbox(label="Transcript (ABAB)", lines=15, value="",
+                          placeholder="Annotations will appear here...")
+    with gr.Row():
+        json_btn = gr.Button("Export JSON")
+        csv_btn  = gr.Button("Export CSV")
+    json_file = gr.File(label="JSON")
+    csv_file  = gr.File(label="CSV")
+
+    run_btn.click(
+        fn=process_audio,
+        inputs=[audio_input, num_speakers],
+        outputs=[status_box, abab_box],
+        queue=False
+    )
+    json_btn.click(fn=export_json, inputs=[], outputs=[json_file])
+    csv_btn.click(fn=export_csv, inputs=[], outputs=[csv_file])
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)