pipeline.py

import os
import json
import csv
import time
import logging
import httpx
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):
    return SPEAKER_LETTERS[idx] if idx < len(SPEAKER_LETTERS) else f"SPK{idx}"

class SpeechAnnotationPipeline:
    def process(self, audio_path, num_speakers=0, **kwargs):
        if not DEEPGRAM_API_KEY:
            raise EnvironmentError("DEEPGRAM_API_KEY not set.")
        logger.info(f"Sending to Deepgram: {audio_path}")
        t0 = time.time()
        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")
        params = {"model":"nova-2","diarize":"true","punctuate":"true",
                  "utterances":"true","smart_format":"true"}
        if num_speakers and num_speakers > 0:
            params["diarize_version"] = "latest"
            params["num_speakers"] = str(num_speakers)
        headers = {"Authorization": f"Token {DEEPGRAM_API_KEY}", "Content-Type": mime}
        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()
        utterances = data.get("results", {}).get("utterances", [])
        if not utterances:
            words = data.get("results",{}).get("channels",[{}])[0]\
                        .get("alternatives",[{}])[0].get("words",[])
            utterances = self._words_to_utterances(words)
        segments = []
        for utt in utterances:
            segments.append(AnnotationSegment(
                speaker=_speaker_label(utt.get("speaker", 0)),
                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):
        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)

_instance = None

def get_pipeline():
    global _instance
    if _instance is None:
        _instance = SpeechAnnotationPipeline()
    return _instance