pipeline_modules/transcriber.py

"""Speech transcription using OpenAI Whisper."""

import logging
import numpy as np
import librosa
from typing import Tuple, Optional

log = logging.getLogger(__name__)

# Cache loaded Whisper models to avoid reloading
_whisper_model_cache = {}


def transcribe_audio(
    signal: np.ndarray,
    sr: int,
    model_name: str = "base",
    language: Optional[str] = None,
) -> Tuple[Optional[str], Optional[str]]:
    """
    Transcribe audio signal using OpenAI Whisper.

    Args:
        signal: Audio signal (1D array)
        sr: Sample rate in Hz
        model_name: Whisper model to use ("tiny", "base", "small", "medium", "large")
        language: Optional language code to hint at (e.g., "en", "da")

    Returns:
        Tuple of (transcription_text, detected_language_code)
        - Returns (None, None) if transcription fails
    """
    try:
        import whisper
    except ImportError:
        log.warning("openai-whisper not installed – skipping transcription")
        return None, None

    try:
        # Load model (with caching to avoid reloading)
        if model_name not in _whisper_model_cache:
            log.info(f"Loading Whisper model '{model_name}'...")
            _whisper_model_cache[model_name] = whisper.load_model(model_name)

        model = _whisper_model_cache[model_name]

        # Resample to 16kHz (Whisper requirement)
        if sr != 16000:
            signal_16k = librosa.resample(signal, orig_sr=sr, target_sr=16000)
        else:
            signal_16k = signal

        signal_16k = signal_16k.astype(np.float32)

        # Transcribe with optional language hint
        kwargs = {"fp16": False}
        if language:
            kwargs["language"] = language

        result = model.transcribe(signal_16k, **kwargs)

        text = result.get("text", "").strip()
        detected_lang = result.get("language", "unknown")

        log.info(f"Transcribed ({detected_lang}): {text[:80]}...")
        return text, detected_lang

    except Exception as e:
        log.warning(f"Whisper transcription failed: {e}")
        return None, None


def transcribe_file(
    wav_path: str,
    model_name: str = "base",
    language: Optional[str] = None,
) -> Tuple[Optional[str], Optional[str]]:
    """
    Transcribe audio file using OpenAI Whisper.

    Args:
        wav_path: Path to WAV file
        model_name: Whisper model to use
        language: Optional language code hint

    Returns:
        Tuple of (transcription_text, detected_language_code)
    """
    try:
        import whisper
    except ImportError:
        log.warning("openai-whisper not installed – skipping transcription")
        return None, None

    try:
        if model_name not in _whisper_model_cache:
            log.info(f"Loading Whisper model '{model_name}'...")
            _whisper_model_cache[model_name] = whisper.load_model(model_name)

        model = _whisper_model_cache[model_name]

        kwargs = {"fp16": False}
        if language:
            kwargs["language"] = language

        result = model.transcribe(wav_path, **kwargs)

        text = result.get("text", "").strip()
        detected_lang = result.get("language", "unknown")

        log.info(f"Transcribed ({detected_lang}): {text[:80]}...")
        return text, detected_lang

    except Exception as e:
        log.warning(f"Whisper transcription failed: {e}")
        return None, None


def clear_model_cache():
    """Clear the Whisper model cache to free memory."""
    global _whisper_model_cache
    _whisper_model_cache.clear()
    log.info("Whisper model cache cleared")