"""The preprocessing contract. Reproduced from pipecat-ai/smart-turn `audio_utils.py` + `inference.py`. Every model in this project -- stock v3.2, E1, the E2 sweep, E3 -- sees audio through THIS module and nothing else. If these functions drift, no number in `results/` is comparable to anything. The order matters and is easy to get wrong: 1. left-zero-pad (or keep the LAST 8 s) to exactly 8 s 2. THEN hand to WhisperFeatureExtractor(chunk_length=8) Doing it the other way round right-pads, and a padding bug measurably hurt official v3.1 before it was fixed in v3.2. """ from __future__ import annotations from functools import lru_cache from pathlib import Path import numpy as np from .config import MAX_AUDIO_S, N_FRAMES, N_MELS, SAMPLE_RATE MAX_SAMPLES = MAX_AUDIO_S * SAMPLE_RATE def truncate_or_leftpad(audio: np.ndarray, n_seconds: int = MAX_AUDIO_S, sample_rate: int = SAMPLE_RATE) -> np.ndarray: """Keep the last `n_seconds`, or left-pad with zeros up to it. Verbatim behaviour of pipecat's `truncate_audio_to_last_n_seconds`. Left padding is deliberate: the decision lives at the END of the clip, so the terminal contour must always land at the same position in the window. """ max_samples = n_seconds * sample_rate if len(audio) > max_samples: return audio[-max_samples:] if len(audio) < max_samples: return np.pad(audio, (max_samples - len(audio), 0), mode="constant", constant_values=0) return audio @lru_cache(maxsize=1) def get_feature_extractor(): from transformers import WhisperFeatureExtractor return WhisperFeatureExtractor(chunk_length=MAX_AUDIO_S) def extract_features(audio: np.ndarray, batched: bool = True) -> np.ndarray: """float32 waveform at 16 kHz -> log-mel (1, 80, 800) float32. `batched=False` returns (80, 800), for writing into a feature cache. """ audio = np.asarray(audio, dtype=np.float32) audio = truncate_or_leftpad(audio) inputs = get_feature_extractor()( audio, sampling_rate=SAMPLE_RATE, return_tensors="np", padding="max_length", max_length=MAX_SAMPLES, truncation=True, do_normalize=True, ) feats = inputs.input_features.squeeze(0).astype(np.float32) return feats[None, ...] if batched else feats def load_audio(path: str | Path, sr: int = SAMPLE_RATE) -> np.ndarray: """Load any file to mono float32 at `sr`, peak-normalised into [-1, 1]. Manifests store repo-relative paths so they survive the trip to Kaggle and the Hub. Resolve against the repo root as a fallback so callers work regardless of the current working directory. """ import librosa p = Path(path) if not p.exists() and not p.is_absolute(): from .config import ROOT if (ROOT / p).exists(): p = ROOT / p audio, _ = librosa.load(str(p), sr=sr, mono=True) audio = audio.astype(np.float32) peak = float(np.max(np.abs(audio))) if audio.size else 0.0 if peak > 1.0: audio = audio / peak return audio def decode_bytes(raw: bytes, sr: int = SAMPLE_RATE) -> np.ndarray: """Decode an in-memory audio blob (FLAC/WAV from HF parquet) to mono 16 kHz.""" import io import librosa import soundfile as sf data, src_sr = sf.read(io.BytesIO(raw), dtype="float32", always_2d=False) if data.ndim > 1: data = data.mean(axis=1) if src_sr != sr: data = librosa.resample(data, orig_sr=src_sr, target_sr=sr) return np.ascontiguousarray(data, dtype=np.float32) def build_ort_session(onnx_path: str | Path): """ONNX Runtime session with pipecat's exact options. These options are part of the latency claim -- benchmarking under different ones would not be comparable to their published 12 ms. """ import onnxruntime as ort so = ort.SessionOptions() so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL so.inter_op_num_threads = 1 so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL return ort.InferenceSession(str(onnx_path), sess_options=so, providers=["CPUExecutionProvider"]) def build_fast_session(onnx_path: str | Path): """Default threading -- for BULK SCORING ONLY. `build_ort_session` deliberately pins ORT_SEQUENTIAL and one inter-op thread because those options ARE Pipecat's published latency contract; benchmarking under anything else would not be comparable. Accuracy scoring has no such contract, so running thousands of clips single-threaded is pure waste. Never use this for latency_bench. """ import onnxruntime as ort so = ort.SessionOptions() so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL return ort.InferenceSession(str(onnx_path), sess_options=so, providers=["CPUExecutionProvider"]) def ort_predict_batch(session, feats: np.ndarray) -> np.ndarray: """(B, 80, 800) -> (B,) probabilities. Uses the graph's dynamic batch axis.""" out = session.run(None, {"input_features": feats.astype(np.float32)}) return np.asarray(out[0]).reshape(-1) def ort_predict(session, feats: np.ndarray) -> float: """Run a session on (1, 80, 800) features. Output is ALREADY sigmoided.""" out = session.run(None, {"input_features": feats.astype(np.float32)}) return float(np.asarray(out[0]).reshape(-1)[0]) def self_test() -> None: """Contract test. Cheap, and it catches the highest-risk bug in the project.""" short = np.ones(SAMPLE_RATE, dtype=np.float32) * 0.5 # 1 s padded = truncate_or_leftpad(short) assert padded.shape == (MAX_SAMPLES,), padded.shape assert np.all(padded[: MAX_SAMPLES - SAMPLE_RATE] == 0), "must pad on the LEFT" assert np.all(padded[MAX_SAMPLES - SAMPLE_RATE:] == 0.5), "signal must land at the END" long = np.arange(12 * SAMPLE_RATE, dtype=np.float32) # 12 s kept = truncate_or_leftpad(long) assert kept.shape == (MAX_SAMPLES,) assert kept[-1] == long[-1], "must keep the LAST 8 s, not the first" feats = extract_features(short) assert feats.shape == (1, N_MELS, N_FRAMES), feats.shape assert feats.dtype == np.float32 assert np.isfinite(feats).all() print(f"audio_utils self-test OK -> {feats.shape} {feats.dtype}") if __name__ == "__main__": self_test()