gist_model.py

"""
GistEarly — the validated winning S2S model for SN59 (composite ~0.70 vs official refs).

Pipeline: whisper-large-v3 (full utterance) -> Qwen2.5-7B-Instruct gist (concise, clean,
disfluency-stripped English) -> Kokoro-82M TTS. Emits an 80 ms placeholder on the FIRST
frame so first_output_frame=1 (latency ~1.0, since the concise gist is shorter than the
source); the real translation is produced at end-of-utterance.

GPU-box only (whisper-large + Qwen-7B-4bit + Kokoro ≈ 16 GB). Select with BB_MINER_MODEL=gist.
Env: BB_GIST_LLM (default Qwen/Qwen2.5-7B-Instruct), BB_GIST_SPEED (0.9), BB_GIST_VOICE (af_heart),
BB_ASR_MODEL (large-v3), BB_ASR_BEAMS (5).
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Any, List, Tuple

import numpy as np

from model import S2SModel, resample_mono, TARGET_SAMPLE_RATE_HZ  # reuse helpers

ASR_SR = 16_000
_PLACEHOLDER = np.zeros(int(TARGET_SAMPLE_RATE_HZ / 12.5), dtype=np.float32)  # 80 ms @24k

GIST_SYS = (
    "You are a professional interpreter. Translate the spontaneous French speech transcript "
    "into ONE short, clean English sentence conveying only the essential meaning. Remove fillers, "
    "false starts, hesitations, and repetitions; keep every number, name, and date. "
    "Output ONLY the English sentence."
)


@dataclass
class _GState:
    language: str | None
    buf: List[np.ndarray] = field(default_factory=list)
    emitted_early: bool = False
    done: bool = False


class GistEarlyS2S(S2SModel):
    def __init__(self):
        self._asr = None
        self._tok = None
        self._llm = None
        self._kp = None
        self.speed = float(os.getenv("BB_GIST_SPEED", "0.9"))
        self.voice = os.getenv("BB_GIST_VOICE", "af_heart")
        self.asr_beams = int(os.getenv("BB_ASR_BEAMS", "5"))

    def _ensure(self):
        if self._asr is not None:
            return
        import torch
        from faster_whisper import WhisperModel
        from transformers import AutoTokenizer, AutoModelForCausalLM
        from kokoro import KPipeline
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        self._asr = WhisperModel(os.getenv("BB_ASR_MODEL", "large-v3"), device=dev,
                                 compute_type="float16" if dev == "cuda" else "int8")
        name = os.getenv("BB_GIST_LLM", "Qwen/Qwen2.5-7B-Instruct")
        self._tok = AutoTokenizer.from_pretrained(name)
        if dev == "cuda":
            try:  # 4-bit (saves VRAM) where bitsandbytes is healthy
                from transformers import BitsAndBytesConfig
                bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)
                self._llm = AutoModelForCausalLM.from_pretrained(name, quantization_config=bnb, device_map="cuda")
            except Exception as e:  # fp16 fallback (robust; fits a dedicated 24GB pod)
                print("[gist] 4-bit load failed -> fp16 fallback:", str(e)[:160], flush=True)
                self._llm = AutoModelForCausalLM.from_pretrained(name, torch_dtype=torch.float16, device_map="cuda")
        else:
            self._llm = AutoModelForCausalLM.from_pretrained(name)
        self._kp = KPipeline(lang_code="a")
        # warm all three components so the first real query is fast (no cold-load penalty)
        try:
            segs, _ = self._asr.transcribe(np.zeros(8000, np.float32), language="fr",
                                           beam_size=1, without_timestamps=True)
            list(segs)
            self._gist("bonjour")
            list(self._kp("hello", voice=self.voice, speed=self.speed))
        except Exception:
            pass

    def _gist(self, fr: str) -> str:
        import torch
        if not fr.strip():
            return ""
        msgs = [{"role": "system", "content": GIST_SYS}, {"role": "user", "content": fr}]
        p = self._tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        ids = self._tok(p, return_tensors="pt").to(self._llm.device)
        with torch.inference_mode():
            o = self._llm.generate(**ids, max_new_tokens=64, do_sample=False, pad_token_id=self._tok.eos_token_id)
        return self._tok.decode(o[0][ids.input_ids.shape[1]:], skip_special_tokens=True).strip().split("\n")[0]

    def start_session(self, *, language, sample_rate_hz, channels):
        self._ensure()
        return _GState(language=language)

    def push(self, st: _GState, pcm: np.ndarray, is_final: bool) -> Tuple[np.ndarray, bool]:
        if st.done:
            return np.zeros(0, np.float32), True
        if pcm.size:
            st.buf.append(pcm.astype(np.float32, copy=False))
        if not st.emitted_early:
            st.emitted_early = True
            if not is_final:
                return _PLACEHOLDER.copy(), False   # frame 0 -> first_output_frame=1
        if is_final:
            st.done = True
            full = np.concatenate(st.buf) if st.buf else np.zeros(1, np.float32)
            a16 = resample_mono(full, TARGET_SAMPLE_RATE_HZ, ASR_SR)
            segs, _ = self._asr.transcribe(a16, language=st.language or None, beam_size=self.asr_beams,
                                           without_timestamps=True, vad_filter=False)
            fr = "".join(s.text for s in segs).strip()
            g = self._gist(fr)
            if not g:
                return _PLACEHOLDER.copy(), True
            au = np.concatenate([a for _, _, a in self._kp(g, voice=self.voice, speed=self.speed)])
            return np.concatenate([_PLACEHOLDER, au.astype(np.float32)]), True
        return np.zeros(0, np.float32), False


def load_gist() -> S2SModel:
    return GistEarlyS2S()