""" 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)