ml-services/src/inference/emotion_predictor.py

"""
Emotion inference pipeline for the audio sentiment analysis module.

This module loads the trained Wav2Vec2 emotion classifier and returns the
standard AudioSentimentResult schema used by the backend and dashboard.

Run test from ml-services:

    python -m src.test_emotion_predictor
"""

import os
from pathlib import Path
from typing import Dict, Optional

import numpy as np
import torch
from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor

from src.data.audio_dataset import (
    DEFAULT_SAMPLE_RATE,
    load_audio_file,
    resolve_audio_path,
)
from src.features.inference_audio_features import (
    extract_audio_feature_summary,
    map_raw_features_to_summary,
)
from src.inference.escalation_score import calculate_audio_escalation_score
from src.inference.sentiment_timeline import (
    TimelineConfig,
    build_sentiment_timeline,
    calculate_audio_sentiment_shift,
    calculate_emotional_volatility,
    find_peak_emotion,
)
from src.sentiment_config import IntensityLevel, SentimentShift
from src.sentiment_schema import (
    AudioSentimentResult,
    EmotionProbabilities,
    PeakEmotion,
    infer_confidence_level,
    infer_overall_sentiment,
    infer_risk_level,
    is_uncertain_prediction,
    seconds_to_timestamp,
)

PROJECT_ROOT = Path(__file__).resolve().parents[3]
ML_SERVICES_ROOT = PROJECT_ROOT / "ml-services"

SELECTED_EMOTION_MODEL_VERSION = "model_v5_cremad_ravdess_freeze6_epochs5_lr1e5"

SELECTED_EMOTION_MODEL_DIR = (
    ML_SERVICES_ROOT
    / "outputs"
    / "wav2vec2"
    / SELECTED_EMOTION_MODEL_VERSION
    / "best_model"
)

DEFAULT_MODEL_DIR = SELECTED_EMOTION_MODEL_DIR


class EmotionPredictor:
    """
    Loads the trained Wav2Vec2 emotion classifier and predicts emotion
    probabilities for one audio file.
    """

    def __init__(
        self,
        model_dir: Path = DEFAULT_MODEL_DIR,
        sample_rate: int = DEFAULT_SAMPLE_RATE,
        max_duration_seconds: Optional[float] = 6.0,
    ) -> None:
        self.model_dir = Path(model_dir)
        self.sample_rate = sample_rate
        self.max_duration_seconds = max_duration_seconds

        if not self.model_dir.exists():
            raise FileNotFoundError(
                f"Selected emotion model not found at: {self.model_dir}\n"
                f"Expected selected model version: {SELECTED_EMOTION_MODEL_VERSION}"
            )

        self.device = self._get_device()

        self.processor = Wav2Vec2Processor.from_pretrained(str(self.model_dir))
        self.model = Wav2Vec2ForSequenceClassification.from_pretrained(str(self.model_dir))

        self.model.to(self.device)
        self.model.eval()

        self.id_to_label = {
            int(index): label for index, label in self.model.config.id2label.items()
        }

    @staticmethod
    def _get_device() -> torch.device:
        """Select the best available device."""
        requested = os.environ.get("ACOUSTIC_DEVICE", "auto").strip().lower()
        if requested != "auto":
            return torch.device(requested)

        if torch.cuda.is_available():
            return torch.device("cuda")

        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return torch.device("mps")

        return torch.device("cpu")

    def predict_probabilities(self, audio_path: Path) -> Dict[str, float]:
        """
        Predict emotion probabilities for one audio file.

        Args:
            audio_path: Path to WAV audio file. Can be relative to ml-services or absolute.

        Returns:
            Dictionary of emotion label to probability.
        """
        resolved_path = resolve_audio_path(str(audio_path))

        waveform, _ = load_audio_file(
            audio_path=resolved_path,
            target_sample_rate=self.sample_rate,
            max_duration_seconds=self.max_duration_seconds,
        )

        inputs = self.processor(
            waveform,
            sampling_rate=self.sample_rate,
            return_tensors="pt",
            padding=True,
            return_attention_mask=True,
        )

        inputs = {key: value.to(self.device) for key, value in inputs.items()}

        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=-1).cpu().numpy()[0]

        emotion_probabilities: Dict[str, float] = {}

        for index, probability in enumerate(probabilities):
            label = self.id_to_label[index]
            emotion_probabilities[label] = float(probability)

        return emotion_probabilities

    @staticmethod
    def _build_probability_schema(
        probabilities: Dict[str, float],
    ) -> EmotionProbabilities:
        """Convert raw model probabilities into the standardized probability schema."""
        return EmotionProbabilities(
            anger=probabilities.get("anger", 0.0),
            disgust=probabilities.get("disgust", 0.0),
            fear=probabilities.get("fear", 0.0),
            happy=probabilities.get("happy", 0.0),
            neutral=probabilities.get("neutral", 0.0),
            sadness=probabilities.get("sadness", 0.0),
        )
    @staticmethod
    def _calculate_confidence_values(
        probabilities: Dict[str, float],
    ) -> tuple[float, float, bool]:
        """
        Calculate confidence, margin, and uncertainty flag.

        Returns:
            prediction_confidence:
                Highest emotion probability.
            top_emotion_margin:
                Difference between top emotion and second emotion.
            uncertain_prediction:
                True if confidence is low or top classes are close.
        """
        if not probabilities:
            return 0.0, 0.0, True

        sorted_probabilities = sorted(
            probabilities.values(),
            reverse=True,
        )

        prediction_confidence = float(sorted_probabilities[0])

        if len(sorted_probabilities) >= 2:
            top_emotion_margin = float(
                sorted_probabilities[0] - sorted_probabilities[1]
            )
        else:
            top_emotion_margin = prediction_confidence

        uncertain = is_uncertain_prediction(
            confidence=prediction_confidence,
            top_emotion_margin=top_emotion_margin,
        )

        return prediction_confidence, top_emotion_margin, uncertain

    @staticmethod
    def _calculate_basic_escalation_score(
        probabilities: EmotionProbabilities,
    ) -> float:
        """
        Calculate first-version escalation score from emotion probabilities only.

        Later steps will add loudness, pitch variability, speech rate, pauses,
        silence duration, and overlap rate.
        """
        score = (
            0.35 * probabilities.anger
            + 0.25 * probabilities.stress_probability()
            + 0.25 * probabilities.negative_probability()
            + 0.15 * probabilities.fear
        )

        return float(np.clip(score, 0.0, 1.0))

    def analyze_audio(
        self,
        audio_path: Path,
        call_id: str = "CALL_TEST_001",
        build_timeline: bool = True,
        raw_audio_features=None,
    ) -> AudioSentimentResult:
        """
        Analyze one audio file and return the official sentiment result.
        """
        raw_probabilities = self.predict_probabilities(audio_path)
        probabilities = self._build_probability_schema(raw_probabilities)
        (
            prediction_confidence,
            top_emotion_margin,
            uncertain_prediction,
        ) = self._calculate_confidence_values(raw_probabilities)

        confidence_level = infer_confidence_level(prediction_confidence)

        dominant_emotion = probabilities.dominant_emotion()
        overall_sentiment = infer_overall_sentiment(probabilities)

        audio_feature_summary = (
            map_raw_features_to_summary(raw_audio_features)
            if raw_audio_features is not None
            else extract_audio_feature_summary(audio_path)
        )

        sentiment_timeline = []
        emotional_volatility = IntensityLevel.LOW
        audio_sentiment_shift = SentimentShift.UNCHANGED
        peak_emotion = PeakEmotion(
            time_seconds=0.0,
            timestamp=seconds_to_timestamp(0.0),
            emotion=dominant_emotion,
            score=self._calculate_basic_escalation_score(probabilities),
        )

        if build_timeline:
            sentiment_timeline = build_sentiment_timeline(
                audio_path=audio_path,
                probability_predictor=self.predict_probabilities,
                config=TimelineConfig(
                    segment_duration_seconds=5.0,
                    min_segment_duration_seconds=1.0,
                    max_duration_seconds=self.max_duration_seconds,
                ),
                single_window_probabilities=raw_probabilities,
            )

            emotional_volatility = calculate_emotional_volatility(sentiment_timeline)
            audio_sentiment_shift = calculate_audio_sentiment_shift(sentiment_timeline)
            peak_emotion = find_peak_emotion(sentiment_timeline)

        escalation_breakdown = calculate_audio_escalation_score(
            probabilities=probabilities,
            audio_features=audio_feature_summary,
            timeline=sentiment_timeline,
            emotional_volatility=emotional_volatility,
            audio_sentiment_shift=audio_sentiment_shift,
            prediction_confidence=prediction_confidence,
            uncertain_prediction=uncertain_prediction,
        )

        escalation_score = escalation_breakdown.final_score
        risk_level = infer_risk_level(escalation_score)

        return AudioSentimentResult(
            call_id=call_id,
            overall_audio_sentiment=overall_sentiment,
            dominant_emotion=dominant_emotion,
            negative_emotion_probability=probabilities.negative_probability(),
            anger_probability=probabilities.anger,
            stress_probability=probabilities.stress_probability(),
            sadness_probability=probabilities.sadness,
            anxiety_probability=probabilities.fear,
            calm_probability=probabilities.calm_probability(),
            audio_features=audio_feature_summary,
            emotional_volatility=emotional_volatility,
            audio_sentiment_shift=audio_sentiment_shift,
            audio_escalation_score=escalation_score,
            risk_level=risk_level,
            escalation_score_breakdown={
                "emotion_risk": escalation_breakdown.emotion_risk,
                "voice_risk": escalation_breakdown.voice_risk,
                "timeline_risk": escalation_breakdown.timeline_risk,
                "uncertainty_adjustment": escalation_breakdown.uncertainty_adjustment,
                "final_score": escalation_breakdown.final_score,
            },
            prediction_confidence=prediction_confidence,
            confidence_level=confidence_level,
            uncertain_prediction=uncertain_prediction,
            top_emotion_margin=top_emotion_margin,
            peak_emotion=peak_emotion,
            sentiment_timeline=sentiment_timeline,
            model_name="Wav2Vec2 emotion classifier - selected V5",
            model_version=SELECTED_EMOTION_MODEL_VERSION,
            processing_status="success",
            warnings=[
                "For short audio clips, the timeline may contain only one segment. For long call-center audio, the same model is applied across multiple segments to track emotion changes over time.",
                "If uncertain_prediction is true, the top emotion probabilities are close or the prediction confidence is low.",
            ],
        )