modeling_dualturn.py

"""modeling_dualturn.py — HuggingFace PreTrainedModel wrapper for DualTurn.

Usage (requires trust_remote_code=True):

    import torch, torchaudio
    from transformers import AutoModel

    model = AutoModel.from_pretrained(
        "anyreach-ai/dualturn-endpointing",
        trust_remote_code=True,
    )
    model.eval()

    wav, sr = torchaudio.load("conversation.wav")  # [2, T]  CH0=user  CH1=agent

    with torch.no_grad():
        out = model(wav, sr=sr)

    # Per-frame signals at 12.5 Hz (one frame = 80 ms)
    print(out.vad_probs.shape)    # [1, T, 2]   (user, agent)
    print(out.eot_probs.shape)    # [1, T, 2]   (user, agent)
    print(out.bot_probs.shape)    # [1, T, 2]   (user, agent)
    print(out.fvad_probs.shape)   # [1, T, 4]   (user_short, user_long, agent_short, agent_long)

    # Endpoint decisions (sparse — only at VAD offsets)
    for ep in out.endpoints:
        print(ep["t_s"], ep["action"], ep["p_st"])
        # ep["action"] is "ST" (user done) or "CL" (user paused)
        # ep["p_st"]   is P(ST) from the endpoint classifier

Notes
-----
- Input audio must be stereo: channel 0 = user, channel 1 = agent.
  Swap channels or set user_channel / agent_channel if needed.
- The model resamples to 24 kHz internally.
- For streaming use the endpointing.py helper (DualTurnEndpointing.stream()).
"""
from __future__ import annotations

import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import joblib
import numpy as np
import torch
import torchaudio
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from configuration_dualturn import DualTurnConfig

# ── Add bundled source to path ────────────────────────────────────────────────
_SRC = Path(__file__).parent / "src"
if str(_SRC) not in sys.path:
    sys.path.insert(0, str(_SRC))


# ─────────────────────────────────────────────────────────────────────────────
# Output dataclass
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class DualTurnOutput(ModelOutput):
    """
    Output from DualTurnModel.

    Attributes
    ----------
    vad_probs : Tensor [batch, T, 2]
        Per-frame voice activity probability.
        dim 2: [user, agent]

    eot_probs : Tensor [batch, T, 2]
        Per-frame end-of-turn probability.
        dim 2: [user, agent]

    bot_probs : Tensor [batch, T, 2]
        Per-frame begin-of-turn probability.
        dim 2: [user, agent]

    fvad_probs : Tensor [batch, T, 4]
        Per-frame fast VAD (Silero) with two smoothing windows per channel.
        dim 2: [user_short, user_long, agent_short, agent_long]

    endpoints : list[dict]
        Sparse endpoint decisions at VAD-offset anchors (agent silent).
        Each dict:
            t_s    : float  — time in seconds
            action : str    — "ST" (user done) or "CL" (user paused)
            p_st   : float  — P(ST) from endpoint classifier  [0, 1]
            signals: dict   — all 10 signal values at the anchor
    """
    vad_probs:  Optional[torch.Tensor] = None   # [B, T, 2]
    eot_probs:  Optional[torch.Tensor] = None   # [B, T, 2]
    bot_probs:  Optional[torch.Tensor] = None   # [B, T, 2]
    fvad_probs: Optional[torch.Tensor] = None   # [B, T, 4]
    endpoints:  list = field(default_factory=list)


# ─────────────────────────────────────────────────────────────────────────────
# Model
# ─────────────────────────────────────────────────────────────────────────────

class DualTurnModel(PreTrainedModel):
    """
    Dual-channel turn-taking model with built-in speech endpoint detection.

    Wraps the dualturn transformer backbone + a trained logistic regression
    endpoint classifier into a single HuggingFace-compatible model.

    Load with:
        model = AutoModel.from_pretrained(
            "anyreach-ai/dualturn-endpointing",
            trust_remote_code=True,
        )
    """

    config_class = DualTurnConfig

    def __init__(self, config: DualTurnConfig):
        super().__init__(config)
        self.config = config

        # Backbone and classifier are loaded lazily from weights in
        # _load_backbone() / _load_classifier(), called by from_pretrained
        self._backbone  = None
        self._input_mode = None
        self._clf       = None
        self._clf_feat_names: list[str] = []
        self._silero_sr  = None

    # ─────────────────────────────────────────────────────────────────────────
    # from_pretrained hook — load backbone weights + classifier
    # ─────────────────────────────────────────────────────────────────────────

    def _init_weights(self, module):
        pass   # weights loaded from checkpoint, not randomly initialised

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        """Override to load backbone (.pt) and classifier (.pkl) from the repo."""
        import os
        from huggingface_hub import snapshot_download

        # Download full repo snapshot to get all files
        local_dir = snapshot_download(pretrained_model_name_or_path,
                                      cache_dir=kwargs.pop("cache_dir", None),
                                      ignore_patterns=["*.md"])

        # Build config
        config = DualTurnConfig.from_pretrained(local_dir)

        # Instantiate with base PreTrainedModel machinery (no weight loading yet)
        kwargs["config"] = config
        model = cls(config)

        # Load dualturn backbone
        backbone_path = os.path.join(local_dir, "best.pt")
        model._load_backbone(backbone_path)

        # Load endpoint classifier
        clf_path = os.path.join(local_dir, "endpoint_clf.pkl")
        model._load_classifier(clf_path)

        return model

    def _load_backbone(self, checkpoint_path: str):
        from evaluation.verify_v2_model import load_model
        from dualturn.data.labeling import silero_probs_native, clean_anti_flicker, SILERO_RATE_HZ

        device = next(iter(["cuda", "cpu"]), "cpu") if not hasattr(self, "device") else str(self.device)
        print(f"[dualturn] Loading backbone from {checkpoint_path} …")
        self._backbone, self._input_mode = load_model(checkpoint_path, device=device)
        self._backbone.eval()

        self._silero_probs   = silero_probs_native
        self._anti_flicker   = clean_anti_flicker
        self._silero_sr      = int(SILERO_RATE_HZ)
        print(f"[dualturn] Backbone ready (input_mode={self._input_mode})")

    def _load_classifier(self, clf_path: str):
        bundle = joblib.load(clf_path)
        self._clf            = bundle["model"]
        self._clf_feat_names = bundle["feature_names"]
        # Honour tuned threshold from bundle unless config overrides
        if "recommended_threshold" in bundle:
            self.config.st_threshold = bundle["recommended_threshold"]
        print(f"[dualturn] Classifier ready "
              f"(type={bundle.get('clf_type','?')}  "
              f"thr={self.config.st_threshold}  "
              f"features={len(self._clf_feat_names)})")

    # ─────────────────────────────────────────────────────────────────────────
    # Forward
    # ─────────────────────────────────────────────────────────────────────────

    def forward(
        self,
        audio: torch.Tensor,
        sr: int = 24_000,
        user_channel: int = 0,
        agent_channel: int = 1,
    ) -> DualTurnOutput:
        """
        Run backbone + endpoint classifier on dual-channel audio.

        Parameters
        ----------
        audio : Tensor [2, T] or [1, 2, T]
            Stereo audio. Channel 0 = user, channel 1 = agent (by default).
        sr : int
            Sample rate of `audio`. Resampled to 24 kHz internally.
        user_channel : int
            Which channel is the user (default 0).
        agent_channel : int
            Which channel is the agent (default 1).

        Returns
        -------
        DualTurnOutput
            .vad_probs  [1, T, 2]   voice activity     (user, agent)
            .eot_probs  [1, T, 2]   end-of-turn        (user, agent)
            .bot_probs  [1, T, 2]   begin-of-turn      (user, agent)
            .fvad_probs [1, T, 4]   fast VAD           (us, ul, as, al)
            .endpoints  list[dict]  sparse ST/CL decisions
        """
        assert self._backbone is not None, "Backbone not loaded. Use from_pretrained()."

        # ── 1. Prepare audio ─────────────────────────────────────────────
        if audio.dim() == 3:
            audio = audio.squeeze(0)   # [2, T]
        if audio.shape[0] != 2:
            raise ValueError(f"Expected 2-channel audio, got shape {audio.shape}")

        mimi_sr = self.config.mimi_sample_rate
        if sr != mimi_sr:
            audio = torchaudio.functional.resample(audio, sr, mimi_sr)

        user_wav  = audio[user_channel].unsqueeze(0)   # [1, T]
        agent_wav = audio[agent_channel].unsqueeze(0)  # [1, T]
        stereo    = torch.stack([audio[user_channel], audio[agent_channel]], dim=0)  # [2, T]

        # ── 2. Backbone inference ─────────────────────────────────────────
        from evaluation.run_model_on_audio import encode_with_mimi, run_inference

        dev = next(self._backbone.parameters()).device
        with torch.no_grad():
            tokens = encode_with_mimi(stereo.unsqueeze(0).to(dev),
                                      self._backbone, self._input_mode)
            preds  = run_inference(tokens, self._backbone)
        # preds: dict of numpy arrays, shape (n_frames,)
        n_frames = preds["vad_user"].shape[0]

        # ── 3. FVAD (Silero) ──────────────────────────────────────────────
        u16 = torchaudio.functional.resample(user_wav,  mimi_sr, self._silero_sr)[0].numpy()
        a16 = torchaudio.functional.resample(agent_wav, mimi_sr, self._silero_sr)[0].numpy()

        raw_u = self._silero_probs(u16, self._silero_sr)
        raw_a = self._silero_probs(a16, self._silero_sr)

        fvad_us = _r2n(self._anti_flicker(raw_u, alpha=self.config.fvad_alpha_short), n_frames)
        fvad_ul = _r2n(self._anti_flicker(raw_u, alpha=self.config.fvad_alpha_long),  n_frames)
        fvad_as = _r2n(self._anti_flicker(raw_a, alpha=self.config.fvad_alpha_short), n_frames)
        fvad_al = _r2n(self._anti_flicker(raw_a, alpha=self.config.fvad_alpha_long),  n_frames)

        # ── 4. Assemble output tensors ────────────────────────────────────
        vad  = torch.tensor(np.stack([preds["vad_user"],  preds["vad_agent"]],  axis=-1)).unsqueeze(0).float()
        eot  = torch.tensor(np.stack([preds["eot_user"],  preds["eot_agent"]],  axis=-1)).unsqueeze(0).float()
        bot  = torch.tensor(np.stack([preds["bot_user"],  preds["bot_agent"]],  axis=-1)).unsqueeze(0).float()
        fvad = torch.tensor(np.stack([fvad_us, fvad_ul, fvad_as, fvad_al],      axis=-1)).unsqueeze(0).float()
        # shapes: [1, T, 2], [1, T, 2], [1, T, 2], [1, T, 4]

        # ── 5. Endpoint decisions ─────────────────────────────────────────
        endpoints = self._detect_endpoints(preds, fvad_us, fvad_ul, fvad_as, fvad_al, n_frames)

        return DualTurnOutput(
            vad_probs  = vad,
            eot_probs  = eot,
            bot_probs  = bot,
            fvad_probs = fvad,
            endpoints  = endpoints,
        )

    # ─────────────────────────────────────────────────────────────────────────
    # Endpoint detection loop
    # ─────────────────────────────────────────────────────────────────────────

    def _detect_endpoints(self, preds, fvad_us, fvad_ul, fvad_as, fvad_al, n_frames):
        """Scan frames for VAD offsets (agent silent) and classify ST/CL."""
        from collections import deque

        cfg          = self.config
        vad_thr      = cfg.vad_edge_threshold
        agent_min    = cfg.agent_voice_min
        ms           = cfg.mimi_frame_ms
        history: deque = deque(maxlen=int(2.0 * cfg.mimi_frame_rate) + 5)

        endpoints = []
        prev_above = False
        last_ep_t  = -999.0

        for i in range(n_frames):
            t_s = i * ms / 1000.0
            sig = {
                "vad_user":         float(preds["vad_user"][i]),
                "vad_agent":        float(preds["vad_agent"][i]),
                "eot_user":         float(preds["eot_user"][i]),
                "eot_agent":        float(preds["eot_agent"][i]),
                "bot_user":         float(preds["bot_user"][i]),
                "bot_agent":        float(preds["bot_agent"][i]),
                "fvad_user_short":  float(fvad_us[i]),
                "fvad_user_long":   float(fvad_ul[i]),
                "fvad_agent_short": float(fvad_as[i]),
                "fvad_agent_long":  float(fvad_al[i]),
            }
            history.append(sig)

            cur_above     = sig["vad_user"] > vad_thr
            recent_agent  = float(np.mean([h["vad_agent"] for h in list(history)[-10:]]))
            agent_voicing = recent_agent > agent_min

            # VAD offset + agent silent + refractory
            if prev_above and not cur_above and not agent_voicing:
                if t_s - last_ep_t >= 0.4:
                    ep = self._classify_endpoint(t_s, sig)
                    endpoints.append(ep)
                    last_ep_t = t_s

            prev_above = cur_above

        return endpoints

    def _classify_endpoint(self, t_s: float, sig: dict) -> dict:
        """
        Run endpoint_clf on the 10 signal values at the VAD offset.

        Features (10, *_last):
          vad_user_last, vad_agent_last, eot_user_last, eot_agent_last,
          bot_user_last, bot_agent_last, fvad_user_short_last,
          fvad_user_long_last, fvad_agent_short_last, fvad_agent_long_last

        Output:
          P(ST) >= threshold → "ST"  (user finished, agent should respond)
          P(ST) <  threshold → "CL"  (user paused mid-sentence, wait)
        """
        SIGNAL_KEYS = [
            "vad_user", "vad_agent", "eot_user", "eot_agent",
            "bot_user", "bot_agent", "fvad_user_short", "fvad_user_long",
            "fvad_agent_short", "fvad_agent_long",
        ]
        feats  = np.array([sig[k] for k in SIGNAL_KEYS], dtype=np.float32).reshape(1, -1)
        proba  = self._clf.predict_proba(feats)[0]
        classes = list(self._clf.classes_)
        p_st   = float(proba[classes.index(1)] if 1 in classes else proba[-1])
        action = "ST" if p_st >= self.config.st_threshold else "CL"
        return {"t_s": t_s, "action": action, "p_st": p_st, "signals": sig}


# ─────────────────────────────────────────────────────────────────────────────
# Helper
# ─────────────────────────────────────────────────────────────────────────────

def _r2n(arr: np.ndarray, n: int) -> np.ndarray:
    """Resample 1-D array to length n via linear interpolation."""
    idx = np.linspace(0, len(arr) - 1, n)
    return np.interp(idx, np.arange(len(arr)), arr).astype(np.float32)