src/audio_utils.py

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