src/models.py

import functools

import numpy as np
import torch

from .config import AST_MODEL, EMOTION_CAT_MODEL, EMOTION_DIM_MODEL, SEGMENTATION_MODEL, SR
from .logger import log

_AUDIOSET_TO_NOISE_TYPE = {
    "Television": "TV",
    "Music": "music",
    "Musical instrument": "music",
    "Vehicle": "traffic/road noise",
    "Car": "traffic/road noise",
    "Traffic noise, roadway noise": "traffic/road noise",
    "Wind": "wind",
    "Wind noise (microphone)": "wind",
    "Typing": "keyboard typing",
    "Computer keyboard": "keyboard typing",
    "Conversation": "office chatter",
    "Chatter": "office chatter",
    "Speech": None,
    "Crowd": "office chatter",
    "Static": "static",
    "White noise": "static",
    "Hum": "mechanical noise",
    "Mechanisms": "mechanical noise",
    "Engine": "mechanical noise",
    "Telephone bell ringing": "office chatter",
    "Dog": "background chatter/animal noise",
    "Silence": None,
}


@functools.lru_cache(maxsize=1)
def _silero_vad():
    model, utils = torch.hub.load(
        "snakers4/silero-vad", "silero_vad", trust_repo=True, onnx=False
    )
    return model, utils


def speech_segments(y: np.ndarray) -> list[tuple[float, float]] | None:
    try:
        model, utils = _silero_vad()
        get_speech_timestamps = utils[0]
        wav = torch.from_numpy(y)
        ts = get_speech_timestamps(wav, model, sampling_rate=SR)
        return [(t["start"] / SR, t["end"] / SR) for t in ts]
    except Exception as e:
        log.exception("speech_segments failed: %s", e)
        return None


@functools.lru_cache(maxsize=1)
def _emotion_model():
    import torch.nn as nn
    from transformers import Wav2Vec2Processor
    from transformers.models.wav2vec2.modeling_wav2vec2 import (
        Wav2Vec2Model,
        Wav2Vec2PreTrainedModel,
    )

    class RegressionHead(nn.Module):
        def __init__(self, config):
            super().__init__()
            self.dense = nn.Linear(config.hidden_size, config.hidden_size)
            self.dropout = nn.Dropout(config.final_dropout)
            self.out_proj = nn.Linear(config.hidden_size, config.num_labels)

        def forward(self, features):
            x = self.dropout(features)
            x = torch.tanh(self.dense(x))
            x = self.dropout(x)
            return self.out_proj(x)

    class EmotionModel(Wav2Vec2PreTrainedModel):
        def __init__(self, config):
            super().__init__(config)
            self.wav2vec2 = Wav2Vec2Model(config)
            self.classifier = RegressionHead(config)
            self.init_weights()

        def forward(self, input_values):
            hidden = self.wav2vec2(input_values)[0]
            pooled = torch.mean(hidden, dim=1)
            return pooled, self.classifier(pooled)

    processor = Wav2Vec2Processor.from_pretrained(EMOTION_DIM_MODEL)
    model = EmotionModel.from_pretrained(EMOTION_DIM_MODEL).eval()
    return processor, model


def dimensional_emotion(y: np.ndarray) -> dict | None:
    try:
        processor, model = _emotion_model()
        inputs = processor(y, sampling_rate=SR, return_tensors="pt")
        with torch.no_grad():
            _, logits = model(inputs["input_values"])
        arousal, dominance, valence = logits[0].tolist()
        return {"arousal": arousal, "dominance": dominance, "valence": valence}
    except Exception as e:
        log.exception("dimensional_emotion failed: %s", e)
        return None


@functools.lru_cache(maxsize=1)
def _categorical_ser_pipeline():
    from transformers import pipeline

    return pipeline("audio-classification", model=EMOTION_CAT_MODEL, top_k=4)


def categorical_emotion(y: np.ndarray) -> list[dict] | None:
    try:
        clf = _categorical_ser_pipeline()
        return clf({"array": y, "sampling_rate": SR})
    except Exception as e:
        log.exception("categorical_emotion failed: %s", e)
        return None


@functools.lru_cache(maxsize=1)
def _ast_pipeline():
    from transformers import pipeline

    return pipeline("audio-classification", model=AST_MODEL, top_k=10)


def noise_tags(y: np.ndarray) -> list[dict] | None:
    try:
        clf = _ast_pipeline()
        return clf({"array": y, "sampling_rate": SR})
    except Exception as e:
        log.exception("noise_tags failed: %s", e)
        return None


def noise_type_from_tags(tags: list[dict], min_score: float = 0.15) -> str:
    for tag in tags:
        mapped = _AUDIOSET_TO_NOISE_TYPE.get(tag["label"])
        if mapped and tag["score"] >= min_score:
            return mapped
    return ""


@functools.lru_cache(maxsize=1)
def _segmentation_inference():
    from pyannote.audio import Inference, Model

    model = Model.from_pretrained(SEGMENTATION_MODEL)
    return Inference(model, step=2.5)


def overlap_seconds(path: str) -> float | None:
    try:
        inference = _segmentation_inference()
        output = inference(path)
        frame_speaker_count = output.data.sum(axis=-1)
        overlap_frames = int((frame_speaker_count >= 2).sum())
        frame_duration = output.sliding_window.duration / output.data.shape[1]
        return float(overlap_frames * frame_duration)
    except Exception as e:
        log.exception("overlap_seconds failed: %s", e)
        return None