streaming_asr.py

"""
Streaming ASR with VAD Endpointing
====================================
Turns push-to-talk into a conversation.

The caller speaks; audio arrives in small chunks. This module decides — with
no button press — when an utterance has STARTED and when it has ENDED, emits
live partial transcripts while the caller is still talking, and detects
barge-in so the agent stops talking when interrupted.

Endpointing state machine
─────────────────────────

     ┌────────┐  speech ≥ min_speech_ms   ┌──────────┐
     │  IDLE  │ ────────────────────────► │ SPEAKING │
     └────────┘                           └──────────┘
         ▲                                      │ silence detected
         │                                      ▼
         │                              ┌────────────────┐
         │  silence ≥ endpoint_ms       │ TRAILING_SIL   │
         └──────── (emit FINAL) ────────│  (may resume)  │
                                        └────────────────┘

Key behaviours:
  - PREROLL      : a ring buffer holds ~300ms of audio from BEFORE the speech
                   trigger fires, so the first phoneme is never clipped. This
                   is the single most common cause of "it dropped my first
                   word" in naive VAD implementations.
  - HANGOVER     : brief silences inside speech (natural pauses between words,
                   the gap before a plosive) do not end the turn. Only
                   `endpoint_silence_ms` of continuous silence does.
  - PARTIALS     : every `partial_interval_ms`, the audio so far is decoded
                   with a SMALL Whisper model for a live on-screen transcript.
                   The FINAL decode uses large-v3 for accuracy.
  - BARGE-IN     : while `agent_speaking` is set, sustained caller speech
                   raises a barge-in event so playback can be cut.
  - MAX DURATION : a hard cap force-endpoints a caller who never pauses.

Whisper hallucinates confidently on silence ("Thank you.", "Subtitles by…"),
so utterances shorter than `min_speech_ms` of actual speech are discarded
without ever reaching the model.
"""

import time
import logging
import numpy as np
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional, Callable

from vad import load_vad, FRAME_SAMPLES, FRAME_MS, SAMPLE_RATE

logger = logging.getLogger(__name__)


class State(Enum):
    IDLE           = "idle"
    SPEAKING       = "speaking"
    TRAILING_SIL   = "trailing_silence"


@dataclass
class EndpointConfig:
    speech_threshold:     float = 0.55    # VAD prob above this = speech frame
    silence_threshold:    float = 0.35    # below this = silence (hysteresis gap
                                          # prevents flapping at the boundary)
    min_speech_ms:        int   = 250     # ignore coughs, door slams, clicks
    endpoint_silence_ms:  int   = 700     # silence that ends a turn
    preroll_ms:           int   = 300     # audio kept from before speech onset
    max_utterance_ms:     int   = 20_000  # hard cap
    partial_interval_ms:  int   = 900     # how often to emit a live partial
    bargein_speech_ms:    int   = 220     # speech needed to interrupt the agent


@dataclass
class UtteranceEvent:
    kind:        str                       # 'partial' | 'final' | 'bargein' |
                                           # 'speech_start' | 'discarded'
    text:        str   = ""
    audio:       Optional[np.ndarray] = None
    duration_ms: float = 0.0
    speech_ms:   float = 0.0
    latency_ms:  float = 0.0


class StreamingASR:
    """
    Feed audio with `accept_audio()`; consume the returned list of events.

    Usage:
        sasr = StreamingASR(transcribe_fn=pipeline.transcribe)
        for chunk in mic_stream:
            for ev in sasr.accept_audio(chunk, sr):
                if ev.kind == "partial": show(ev.text)
                if ev.kind == "final":   handle(ev.text)
    """

    def __init__(self,
                 transcribe_fn: Callable[[np.ndarray, int], str],
                 partial_transcribe_fn: Optional[Callable] = None,
                 config: Optional[EndpointConfig] = None,
                 vad_backend: str = "auto",
                 emit_partials: bool = True):
        self.cfg = config or EndpointConfig()
        self.vad = load_vad(vad_backend)
        self.transcribe_fn = transcribe_fn
        # Partials can use a smaller/faster model; falls back to the main one
        self.partial_transcribe_fn = partial_transcribe_fn or transcribe_fn
        self.emit_partials = emit_partials

        self.agent_speaking = False        # set True while TTS plays (barge-in)
        self._preroll_frames = max(1, int(self.cfg.preroll_ms / FRAME_MS))
        self.reset()

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    def reset(self):
        self.state           = State.IDLE
        self._buffer         = np.zeros(0, dtype=np.float32)   # leftover samples
        self._preroll        = []                              # ring of frames
        self._utterance      = []                              # frames of turn
        self._speech_ms      = 0.0
        self._silence_ms     = 0.0
        self._utterance_ms   = 0.0
        self._bargein_ms     = 0.0
        self._last_partial_ms = 0.0
        self._partial_text   = ""
        self.vad.reset()

    # ── Main entry ────────────────────────────────────────────────────────────

    def accept_audio(self, audio: np.ndarray,
                     sample_rate: int = SAMPLE_RATE) -> list[UtteranceEvent]:
        """
        audio: mono float32 in [-1,1] (int16 is auto-converted) of any length.
        Returns zero or more events produced by this chunk.
        """
        events: list[UtteranceEvent] = []
        audio = _to_float_mono(audio)
        if sample_rate != SAMPLE_RATE:
            audio = _resample(audio, sample_rate, SAMPLE_RATE)

        self._buffer = np.concatenate([self._buffer, audio])

        # Consume whole frames only; remainder stays buffered for next chunk
        while len(self._buffer) >= FRAME_SAMPLES:
            frame = self._buffer[:FRAME_SAMPLES]
            self._buffer = self._buffer[FRAME_SAMPLES:]
            ev = self._process_frame(frame)
            events.extend(ev)

        return events

    def flush(self) -> list[UtteranceEvent]:
        """Force-endpoint whatever is buffered (e.g. caller hung up)."""
        if self.state in (State.SPEAKING, State.TRAILING_SIL):
            return self._finalize()
        return []

    # ── Frame processing ──────────────────────────────────────────────────────

    def _process_frame(self, frame: np.ndarray) -> list[UtteranceEvent]:
        events = []
        prob = self.vad.speech_prob(frame)
        is_speech  = prob >= self.cfg.speech_threshold
        is_silence = prob <= self.cfg.silence_threshold

        # ── Barge-in: caller talks over the agent ────────────────────────────
        if self.agent_speaking:
            self._bargein_ms = self._bargein_ms + FRAME_MS if is_speech else 0.0
            if self._bargein_ms >= self.cfg.bargein_speech_ms:
                self._bargein_ms = 0.0
                self.agent_speaking = False
                events.append(UtteranceEvent(kind="bargein"))
                # fall through — this frame also starts the new utterance

        # ── Preroll ring (only meaningful while IDLE) ────────────────────────
        if self.state is State.IDLE:
            self._preroll.append(frame)
            if len(self._preroll) > self._preroll_frames:
                self._preroll.pop(0)

        # ── State machine ────────────────────────────────────────────────────
        if self.state is State.IDLE:
            if is_speech:
                self._speech_ms += FRAME_MS
                if self._speech_ms >= self.cfg.min_speech_ms:
                    # Commit: open the utterance with the preroll in front
                    self._utterance = list(self._preroll)
                    self._utterance_ms = len(self._utterance) * FRAME_MS
                    self._preroll = []
                    self._silence_ms = 0.0
                    self._last_partial_ms = 0.0
                    self.state = State.SPEAKING
                    events.append(UtteranceEvent(kind="speech_start"))
            else:
                self._speech_ms = 0.0
            return events

        # SPEAKING or TRAILING_SIL — always accumulate audio
        self._utterance.append(frame)
        self._utterance_ms += FRAME_MS

        if self.state is State.SPEAKING:
            if is_silence:
                self.state = State.TRAILING_SIL
                self._silence_ms = FRAME_MS
            else:
                if is_speech:
                    self._speech_ms += FRAME_MS
                self._silence_ms = 0.0

        elif self.state is State.TRAILING_SIL:
            if is_speech:
                # Natural pause, not an endpoint — resume
                self.state = State.SPEAKING
                self._speech_ms += FRAME_MS
                self._silence_ms = 0.0
            else:
                self._silence_ms += FRAME_MS
                if self._silence_ms >= self.cfg.endpoint_silence_ms:
                    return events + self._finalize()

        # Hard cap on a caller who never pauses
        if self._utterance_ms >= self.cfg.max_utterance_ms:
            logger.info("Max utterance length reached — force endpoint.")
            return events + self._finalize()

        # ── Live partial transcript ──────────────────────────────────────────
        if (self.emit_partials
                and self.state is State.SPEAKING
                and self._utterance_ms - self._last_partial_ms
                    >= self.cfg.partial_interval_ms):
            self._last_partial_ms = self._utterance_ms
            ev = self._emit_partial()
            if ev:
                events.append(ev)

        return events

    # ── Emission ──────────────────────────────────────────────────────────────

    def _emit_partial(self) -> Optional[UtteranceEvent]:
        audio = np.concatenate(self._utterance)
        t0 = time.perf_counter()
        try:
            text = self.partial_transcribe_fn(audio, SAMPLE_RATE)
        except Exception as e:
            logger.warning(f"Partial decode failed: {e}")
            return None
        text = (text or "").strip()
        if not text or text == self._partial_text:
            return None
        self._partial_text = text
        return UtteranceEvent(
            kind="partial", text=text,
            duration_ms=self._utterance_ms, speech_ms=self._speech_ms,
            latency_ms=(time.perf_counter() - t0) * 1000)

    def _finalize(self) -> list[UtteranceEvent]:
        audio      = np.concatenate(self._utterance) if self._utterance else np.zeros(0)
        speech_ms  = self._speech_ms
        total_ms   = self._utterance_ms

        # Reset BEFORE decoding so late-arriving audio starts a clean turn
        self._utterance   = []
        self._preroll     = []
        self._speech_ms   = 0.0
        self._silence_ms  = 0.0
        self._utterance_ms = 0.0
        self._partial_text = ""
        self.state = State.IDLE
        self.vad.reset()

        # Guard: never send near-silence to Whisper (hallucination source)
        if speech_ms < self.cfg.min_speech_ms or len(audio) < FRAME_SAMPLES * 4:
            logger.info(f"Discarded short utterance ({speech_ms:.0f}ms speech).")
            return [UtteranceEvent(kind="discarded", speech_ms=speech_ms,
                                    duration_ms=total_ms)]

        t0 = time.perf_counter()
        try:
            text = self.transcribe_fn(audio, SAMPLE_RATE)
        except Exception as e:
            logger.error(f"Final decode failed: {e}")
            return [UtteranceEvent(kind="discarded", speech_ms=speech_ms,
                                    duration_ms=total_ms)]
        latency = (time.perf_counter() - t0) * 1000

        text = (text or "").strip()
        if not text or _is_hallucination(text):
            logger.info(f"Discarded empty/hallucinated final: {text!r}")
            return [UtteranceEvent(kind="discarded", speech_ms=speech_ms,
                                    duration_ms=total_ms)]

        logger.info(f"FINAL ({latency:.0f}ms, {speech_ms:.0f}ms speech): {text}")
        return [UtteranceEvent(kind="final", text=text, audio=audio,
                                duration_ms=total_ms, speech_ms=speech_ms,
                                latency_ms=latency)]


# ── Hallucination filter ──────────────────────────────────────────────────────

_HALLUCINATIONS = {
    "thank you.", "thanks for watching!", "thank you for watching.",
    "you", ".", "...", "subtitles by the amara.org community",
    "please subscribe.", "bye.", "amara.org", "sous-titrage",
    "merci d'avoir regardé cette vidéo!", "à suivre",
}


def _is_hallucination(text: str) -> bool:
    t = text.strip().lower()
    if t in _HALLUCINATIONS:
        return True
    # A "sentence" of only punctuation / music tags
    if all(c in " .,!?-–—♪[]()" for c in t):
        return True
    return False


# ── Audio helpers ─────────────────────────────────────────────────────────────

def _to_float_mono(audio: np.ndarray) -> np.ndarray:
    audio = np.asarray(audio)
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    if audio.dtype == np.int16:
        audio = audio.astype(np.float32) / 32768.0
    elif audio.dtype == np.int32:
        audio = audio.astype(np.float32) / 2147483648.0
    else:
        audio = audio.astype(np.float32)
    return audio


def _resample(audio: np.ndarray, src: int, dst: int) -> np.ndarray:
    if src == dst:
        return audio
    try:
        import scipy.signal as ss
        n = int(round(len(audio) * dst / src))
        return ss.resample(audio, n).astype(np.float32)
    except Exception:
        # Linear interpolation fallback
        n = int(round(len(audio) * dst / src))
        xp = np.linspace(0, 1, len(audio), endpoint=False)
        x  = np.linspace(0, 1, n, endpoint=False)
        return np.interp(x, xp, audio).astype(np.float32)