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