""" KES Transcription API — HuggingFace Space (Free CPU tier). Exposes a POST /transcribe endpoint that accepts audio via URL or base64, runs Whisper + optional pyannote diarization, and returns the rich JSON shape expected by the transcription microservice. """ import base64 import os import tempfile import warnings from typing import Any, Dict, List, Optional import requests from fastapi import FastAPI, HTTPException from pydantic import BaseModel warnings.filterwarnings("ignore") # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- WHISPER_MODEL = os.environ.get("WHISPER_MODEL", "small") HF_TOKEN = os.environ.get("HF_TOKEN", "") ENABLE_DIARIZATION = os.environ.get("ENABLE_DIARIZATION", "false").lower() == "true" # --------------------------------------------------------------------------- # Global model references (loaded once at startup) # --------------------------------------------------------------------------- whisper_model = None diarization_pipeline = None def load_whisper(): global whisper_model from faster_whisper import WhisperModel import torch device = "cuda" if torch.cuda.is_available() else "cpu" compute_type = "float16" if device == "cuda" else "int8" print(f"Loading Whisper '{WHISPER_MODEL}' on {device} ({compute_type})...") whisper_model = WhisperModel(WHISPER_MODEL, device=device, compute_type=compute_type) print("Whisper loaded.") def patch_torchaudio(): """Patch torchaudio for compatibility with pyannote 3.x on newer PyTorch.""" import torchaudio if not hasattr(torchaudio, "AudioMetaData"): from dataclasses import dataclass @dataclass class AudioMetaData: sample_rate: int = 0 num_frames: int = 0 num_channels: int = 0 bits_per_sample: int = 0 encoding: str = "" torchaudio.AudioMetaData = AudioMetaData if not hasattr(torchaudio, "info"): import soundfile as sf def _info(filepath): info = sf.info(filepath) return torchaudio.AudioMetaData( sample_rate=info.samplerate, num_frames=info.frames, num_channels=info.channels, bits_per_sample=16, encoding=info.subtype or "PCM_16", ) torchaudio.info = _info if not hasattr(torchaudio, "list_audio_backends"): torchaudio.list_audio_backends = lambda: ["soundfile"] def load_diarization(): global diarization_pipeline if not ENABLE_DIARIZATION: print("Diarization disabled (set ENABLE_DIARIZATION=true to enable).") return if not HF_TOKEN: print("Warning: HF_TOKEN not set, diarization disabled.") return try: patch_torchaudio() import torch from pyannote.audio import Pipeline os.environ["HF_TOKEN"] = HF_TOKEN print("Loading pyannote diarization pipeline...") diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1") print("Diarization loaded.") except Exception as exc: print(f"Warning: diarization disabled: {exc}") diarization_pipeline = None # --------------------------------------------------------------------------- # FastAPI app # --------------------------------------------------------------------------- app = FastAPI(title="KES Transcription API") class TranscribeRequest(BaseModel): audio_url: Optional[str] = None audio_base64: Optional[str] = None filename: Optional[str] = None task: Optional[str] = None language: Optional[str] = None num_speakers: int = 2 class WordItem(BaseModel): text: str start: float end: float type: str = "word" speaker_id: str = "SPEAKER_00" confidence: float = 0.0 class TranscribeResponse(BaseModel): text: str language: str language_probability: float duration: float avg_confidence: float words: List[WordItem] speakers: List[str] @app.on_event("startup") def startup(): load_whisper() load_diarization() @app.post("/") async def root_transcribe(body: dict): """Accept Inference Endpoint format for compatibility with the microservice. Expected body: { "inputs": {"audio_url": "..."} or {"audio_base64": "...", "filename": "..."}, "parameters": {"task": "translate", "language": "en", "num_speakers": 2} } """ inputs = body.get("inputs", body) parameters = body.get("parameters") or {} req = TranscribeRequest( audio_url=inputs.get("audio_url") if isinstance(inputs, dict) else None, audio_base64=inputs.get("audio_base64") if isinstance(inputs, dict) else None, filename=inputs.get("filename") if isinstance(inputs, dict) else None, task=parameters.get("task"), language=parameters.get("language"), num_speakers=int(parameters.get("num_speakers") or 2), ) return transcribe(req) @app.get("/health") def health(): return { "status": "ok", "whisper_model": WHISPER_MODEL, "diarization": diarization_pipeline is not None, } @app.post("/transcribe", response_model=TranscribeResponse) def transcribe(req: TranscribeRequest): if not req.audio_url and not req.audio_base64: raise HTTPException(status_code=400, detail="Provide audio_url or audio_base64") # Materialise audio to a temp file audio_path = materialise_audio(req) try: # Transcribe transcription = run_transcription(audio_path, req.language, req.task) # Diarize (if enabled) speaker_segments = run_diarization(audio_path, req.num_speakers) # Merge words = merge_words_with_speakers(transcription["words"], speaker_segments) speakers = sorted({w["speaker_id"] for w in words}) if words else ["SPEAKER_00"] avg_conf = 0.0 confs = [w["confidence"] for w in words] if confs: avg_conf = round(sum(confs) / len(confs), 4) return TranscribeResponse( text=transcription["text"], language=transcription["language"], language_probability=transcription["language_probability"], duration=transcription["duration"], avg_confidence=avg_conf, words=[WordItem(**w) for w in words], speakers=speakers, ) finally: try: os.unlink(audio_path) except OSError: pass # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def materialise_audio(req: TranscribeRequest) -> str: suffix = ".wav" if req.audio_url: resp = requests.get(req.audio_url, timeout=300) resp.raise_for_status() audio_bytes = resp.content clean = req.audio_url.split("?", 1)[0] ext = os.path.splitext(clean)[1] if ext: suffix = ext else: audio_bytes = base64.b64decode(req.audio_base64) if req.filename: ext = os.path.splitext(req.filename)[1] if ext: suffix = ext handle = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) try: handle.write(audio_bytes) return handle.name finally: handle.close() def run_transcription( audio_path: str, language: Optional[str], task: Optional[str] ) -> Dict[str, Any]: kwargs: Dict[str, Any] = {"word_timestamps": True, "vad_filter": True} if language: kwargs["language"] = language if task in {"transcribe", "translate"}: kwargs["task"] = task segments_gen, info = whisper_model.transcribe(audio_path, **kwargs) words: List[Dict[str, Any]] = [] text_parts: List[str] = [] for segment in segments_gen: text_parts.append(segment.text.strip()) for word in segment.words or []: words.append( { "text": word.word, "start": round(word.start, 3), "end": round(word.end, 3), "type": "word", "speaker_id": "SPEAKER_00", "confidence": round(word.probability, 4), } ) return { "text": " ".join(text_parts).strip(), "language": info.language, "language_probability": round(info.language_probability, 4), "duration": round(info.duration, 2), "words": words, } def run_diarization(audio_path: str, num_speakers: int) -> Optional[List[Dict[str, Any]]]: if diarization_pipeline is None: return None try: import soundfile as sf import torch data, sample_rate = sf.read(audio_path) if data.ndim == 1: waveform = torch.from_numpy(data).float().unsqueeze(0) else: waveform = torch.from_numpy(data.T).float() diarization = diarization_pipeline( {"waveform": waveform, "sample_rate": sample_rate}, num_speakers=num_speakers, ) segments: List[Dict[str, Any]] = [] for turn, _, speaker in diarization.itertracks(yield_label=True): segments.append( { "start": round(turn.start, 3), "end": round(turn.end, 3), "speaker": speaker, } ) return segments except Exception as exc: print(f"Warning: diarization failed: {exc}") return None def merge_words_with_speakers( words: List[Dict[str, Any]], speaker_segments: Optional[List[Dict[str, Any]]], ) -> List[Dict[str, Any]]: if not speaker_segments: return words for word in words: midpoint = (word["start"] + word["end"]) / 2 best_speaker = "SPEAKER_00" best_overlap = -1.0 for seg in speaker_segments: if seg["start"] <= midpoint <= seg["end"]: overlap = min(word["end"], seg["end"]) - max(word["start"], seg["start"]) if overlap > best_overlap: best_overlap = overlap best_speaker = seg["speaker"] word["speaker_id"] = best_speaker return words