pipeline.py

"""
HausaVoiceAIPipeline
====================
Full pipeline: Audio → Whisper ASR → NLLB translation → LLM NLU → NLLB → MMS-TTS → Audio

Models used:
  - ASR  : openai/whisper-large-v3          (Hausa language support)
  - MT   : facebook/nllb-200-distilled-600M  (hau_Latn ↔ eng_Latn)
  - TTS  : facebook/mms-tts-hau              (Hausa VITS synthesis)
"""

import os
import time
import logging
import numpy as np
import torch
from typing import Optional, Tuple

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ── Language codes ──────────────────────────────────────────────────────────
HAUSA_NLLB   = "hau_Latn"
ENGLISH_NLLB = "eng_Latn"
FRENCH_NLLB  = "fra_Latn"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"Pipeline running on: {DEVICE}")


class HausaVoiceAIPipeline:
    """
    Lazy-loading pipeline.  Each component is loaded on first use to keep
    the Space startup time reasonable on CPU.
    """

    # Model selection is CPU-FIRST. whisper-large-v3 is ~6GB in fp32 and takes
    # roughly a minute per utterance on the 2-vCPU free tier — it makes the
    # Space look broken. Default to small on CPU, large-v3 when a GPU is
    # present, and let an env var override either way.
    _GPU = torch.cuda.is_available()
    ASR_FINAL_MODEL = os.getenv(
        "ASR_FINAL_MODEL",
        "openai/whisper-large-v3" if _GPU else "openai/whisper-small")
    ASR_PARTIAL_MODEL = os.getenv(
        "ASR_PARTIAL_MODEL",
        "openai/whisper-small" if _GPU else "openai/whisper-base")
    NLLB_MODEL = os.getenv("NLLB_MODEL", "facebook/nllb-200-distilled-600M")
    TTS_MODEL  = os.getenv("TTS_MODEL",  "facebook/mms-tts-hau")

    def describe_models(self) -> str:
        short = lambda m: m.split("/")[-1]
        return (f"{short(self.ASR_FINAL_MODEL)} · {short(self.NLLB_MODEL)} · "
                f"{short(self.TTS_MODEL)} · {'GPU' if self._GPU else 'CPU'}")

    def __init__(self, pivot_language: str = "english"):
        self.pivot = pivot_language          # dialogue logic runs in English
        self._asr   = None
        self._asr_fast = None
        self._nllb_model     = None
        self._nllb_tokenizer = None
        self._tts_model      = None
        self._tts_tokenizer  = None
        self.sample_rate = 16_000           # Whisper input
        self.tts_sample_rate = 16_000       # MMS output

    # ── Lazy loaders ────────────────────────────────────────────────────────

    def _load_asr(self):
        if self._asr is not None:
            return
        logger.info(f"Loading {self.ASR_FINAL_MODEL} …")
        from transformers import pipeline as hf_pipeline
        self._asr = hf_pipeline(
            "automatic-speech-recognition",
            model=self.ASR_FINAL_MODEL,
            generate_kwargs={"language": "hausa", "task": "transcribe"},
            device=0 if DEVICE == "cuda" else -1,
            chunk_length_s=30,
        )

    def _load_asr_fast(self):
        """Small model for live partial transcripts."""
        if self._asr_fast is not None:
            return
        if self.ASR_PARTIAL_MODEL == self.ASR_FINAL_MODEL:
            self._load_asr()
            self._asr_fast = self._asr
            return
        logger.info(f"Loading {self.ASR_PARTIAL_MODEL} (partials) …")
        from transformers import pipeline as hf_pipeline
        self._asr_fast = hf_pipeline(
            "automatic-speech-recognition",
            model=self.ASR_PARTIAL_MODEL,
            generate_kwargs={"language": "hausa", "task": "transcribe"},
            device=0 if DEVICE == "cuda" else -1,
            chunk_length_s=30,
        )

    def _load_nllb(self):
        if self._nllb_model is not None:
            return
        logger.info("Loading NLLB-200-distilled-600M …")
        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
        model_id = self.NLLB_MODEL
        self._nllb_tokenizer = AutoTokenizer.from_pretrained(model_id)
        self._nllb_model = AutoModelForSeq2SeqLM.from_pretrained(model_id).to(DEVICE)

    def _load_tts(self):
        if self._tts_model is not None:
            return
        logger.info("Loading MMS-TTS Hausa …")
        from transformers import VitsModel, AutoTokenizer
        model_id = self.TTS_MODEL
        self._tts_tokenizer = AutoTokenizer.from_pretrained(model_id)
        self._tts_model = VitsModel.from_pretrained(model_id).to(DEVICE)
        self.tts_sample_rate = self._tts_model.config.sampling_rate

    # ── Public API ───────────────────────────────────────────────────────────

    def transcribe(self, audio_array: np.ndarray, sample_rate: int = 16_000) -> str:
        """
        Whisper ASR: raw audio → Hausa text.
        audio_array : float32 numpy array, mono
        """
        self._load_asr()
        t0 = time.perf_counter()
        # Resample if needed (scipy-based; avoids pulling in librosa/numba)
        if sample_rate != 16_000:
            from streaming_asr import _resample
            audio_array = _resample(audio_array, sample_rate, 16_000)
        result = self._asr({"array": audio_array.astype(np.float32), "sampling_rate": 16_000})
        text = result["text"].strip()
        logger.info(f"ASR ({time.perf_counter()-t0:.2f}s): {text}")
        return text

    def transcribe_partial(self, audio_array: np.ndarray,
                           sample_rate: int = 16_000) -> str:
        """Fast, lower-accuracy decode for live on-screen partials."""
        self._load_asr_fast()
        if sample_rate != 16_000:
            from streaming_asr import _resample
            audio_array = _resample(audio_array, sample_rate, 16_000)
        result = self._asr_fast({"array": audio_array.astype(np.float32),
                                  "sampling_rate": 16_000})
        return result["text"].strip()

    def make_streaming_session(self, emit_partials: bool = True,
                                config=None, vad_backend: str = "auto"):
        """
        Build a StreamingASR wired to this pipeline's two Whisper models.
        Each caller/session needs its own instance (it holds audio state).
        """
        from streaming_asr import StreamingASR
        return StreamingASR(
            transcribe_fn=self.transcribe,
            partial_transcribe_fn=self.transcribe_partial,
            config=config,
            vad_backend=vad_backend,
            emit_partials=emit_partials,
        )

    def translate(self, text: str, src_lang: str, tgt_lang: str,
                  max_new_tokens: int = 256) -> str:
        """NLLB translation."""
        self._load_nllb()
        t0 = time.perf_counter()
        self._nllb_tokenizer.src_lang = src_lang
        inputs = self._nllb_tokenizer(text, return_tensors="pt", truncation=True,
                                       max_length=512).to(DEVICE)
        forced_id = self._nllb_tokenizer.convert_tokens_to_ids(tgt_lang)
        with torch.no_grad():
            tokens = self._nllb_model.generate(
                **inputs,
                forced_bos_token_id=forced_id,
                max_new_tokens=max_new_tokens,
            )
        translated = self._nllb_tokenizer.decode(tokens[0], skip_special_tokens=True)
        logger.info(f"NLLB {src_lang}→{tgt_lang} ({time.perf_counter()-t0:.2f}s): {translated}")
        return translated

    def synthesize(self, text: str) -> Tuple[int, np.ndarray]:
        """
        MMS-TTS: Hausa text → (sample_rate, audio_array int16).
        """
        self._load_tts()
        t0 = time.perf_counter()
        inputs = self._tts_tokenizer(text, return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            output = self._tts_model(**inputs).waveform
        audio = output.squeeze().cpu().numpy()
        # Normalise → int16
        audio = (audio / np.abs(audio).max() * 32767).astype(np.int16)
        logger.info(f"TTS ({time.perf_counter()-t0:.2f}s): {len(audio)/self.tts_sample_rate:.1f}s audio")
        return self.tts_sample_rate, audio

    # ── Full round-trip helper ───────────────────────────────────────────────

    def hausa_to_english(self, hausa_text: str) -> str:
        return self.translate(hausa_text, HAUSA_NLLB, ENGLISH_NLLB)

    def english_to_hausa(self, english_text: str) -> str:
        return self.translate(english_text, ENGLISH_NLLB, HAUSA_NLLB)

    def audio_to_hausa_text(self, audio_array: np.ndarray,
                             sample_rate: int = 16_000) -> str:
        return self.transcribe(audio_array, sample_rate)

    def hausa_text_to_audio(self, hausa_text: str) -> Tuple[int, np.ndarray]:
        return self.synthesize(hausa_text)