Real Wav2Vec2 SER model inference from model.safetensors

app.py

@@ -7,8 +7,13 @@ import os
 from typing import Optional
 
 import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
 from fastapi import FastAPI, File, Form, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from safetensors.torch import load_file
+from transformers import Wav2Vec2Model, Wav2Vec2Config
 
 app = FastAPI(title="Speech Emotion Recognition")
 
@@ -20,23 +25,47 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ── Model output classes (7-class SER model) ──────────────────────────
 MODEL_CLASSES = ["angry", "disgust", "fear", "happy", "neutral", "pleasant_surprise", "sad"]
 MODEL_CLASS_INDEX = {c: i for i, c in enumerate(MODEL_CLASSES)}
 
-# ── Model ─────────────────────────────────────────────────────────────
 _model = None
 
 
+class SERHead(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.projector = nn.Linear(768, 256)
+        self.classifier = nn.Linear(256, 7)
+        self.layer_weights = nn.Parameter(torch.ones(13) / 13)
+
+    def forward(self, hidden_states):
+        stacked = torch.stack(list(hidden_states), dim=0)
+        w = F.softmax(self.layer_weights, dim=0)
+        weighted = (stacked * w.view(-1, 1, 1, 1)).sum(dim=0)
+        pooled = weighted.mean(dim=1)
+        return self.classifier(F.relu(self.projector(pooled)))
+
+
 def load_model():
-    """Load the speech emotion model. Called once at startup."""
     global _model
-    # TODO: Replace with actual model loading, e.g. PyTorch:
-    #   import torch
-    #   _model = torch.jit.load("model.pt")
-    #   _model.eval()
-    print("[INFO] No model loaded — using dummy predictor")
-    _model = "dummy"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    config = Wav2Vec2Config.from_pretrained("facebook/wav2vec2-base")
+    backbone = Wav2Vec2Model(config).to(device).eval()
+    head = SERHead().to(device).eval()
+
+    model_path = "model.safetensors"
+    if not os.path.exists(model_path):
+        print(f"[WARN] {model_path} not found — using dummy predictor")
+        _model = "dummy"
+        return
+
+    state = load_file(model_path)
+    backbone.load_state_dict(state, strict=False)
+    head.load_state_dict(state, strict=False)
+
+    _model = {"backbone": backbone, "head": head, "device": device}
+    print(f"[INFO] Model loaded on {device}")
 
 
 @app.on_event("startup")
@@ -45,48 +74,42 @@ async def startup():
 
 
 def _extract_features(waveform: np.ndarray, sr: int) -> np.ndarray:
-    """Extract MFCC features for model input."""
     import librosa
     mfcc = librosa.feature.mfcc(y=waveform, sr=sr, n_mfcc=40, n_fft=512, hop_length=256)
     mfcc_delta = librosa.feature.delta(mfcc)
     mfcc_delta2 = librosa.feature.delta(mfcc, order=2)
     features = np.concatenate([mfcc, mfcc_delta, mfcc_delta2], axis=0)
-    # Mean-pool over time → fixed-size vector
     return features.mean(axis=1)
 
 
 def _predict(waveform: np.ndarray, sr: int) -> dict:
-    """
-    Run the loaded model on waveform.
-    Replace this function with actual inference when _model is a real model.
-    """
+    global _model
     if _model is None or _model == "dummy":
         return _dummy_predict(waveform, sr)
 
-    # ── TODO: Replace with real model inference ───────────────────────
-    # Example for a PyTorch model that takes MFCC features:
-    #
-    #   feats = _extract_features(waveform, sr)       # shape: (120,)
-    #   import torch
-    #   with torch.no_grad():
-    #       logits = _model(torch.from_numpy(feats).unsqueeze(0))
-    #       probs = torch.softmax(logits, dim=-1).squeeze(0).numpy()
-    #
-    # For a TensorFlow/Keras model:
-    #
-    #   feats = _extract_features(waveform, sr).reshape(1, -1)
-    #   probs = _model.predict(feats, verbose=0)[0]
-    #
-    # Then construct response:
-    #   pred_idx = int(probs.argmax())
-    #   emotion = MODEL_CLASSES[pred_idx]
-    #   prob_map = {c: round(float(probs[i]), 4) for i, c in enumerate(MODEL_CLASSES)}
-
-    return _dummy_predict(waveform, sr)
+    backbone = _model["backbone"]
+    head = _model["head"]
+    device = _model["device"]
+
+    wav_t = torch.from_numpy(waveform).float().unsqueeze(0).to(device)
+    with torch.no_grad():
+        outputs = backbone(wav_t, output_hidden_states=True)
+        logits = head(outputs.hidden_states)
+        probs = F.softmax(logits, dim=-1).squeeze(0)
+
+    probs_np = probs.cpu().numpy()
+    pred_idx = int(probs_np.argmax())
+    emotion = MODEL_CLASSES[pred_idx]
+    prob_map = {c: round(float(probs_np[i]), 4) for i, c in enumerate(MODEL_CLASSES)}
+
+    return {
+        "emotion": emotion,
+        "confidence": round(float(probs_np[pred_idx]), 4),
+        "probabilities": prob_map,
+    }
 
 
 def _dummy_predict(waveform: np.ndarray, sr: int) -> dict:
-    """Dummy predictor — returns neutral with random noise."""
     rng = np.random.default_rng(42)
     probs = rng.dirichlet(np.ones(len(MODEL_CLASSES)) * 0.5)
     probs = probs / probs.sum()
@@ -99,7 +122,6 @@ def _dummy_predict(waveform: np.ndarray, sr: int) -> dict:
 
 
 def _predict_from_bytes(wav_bytes: bytes) -> dict:
-    """Convert raw WAV bytes to features and run the model."""
     try:
         import soundfile as sf
         import librosa
@@ -118,20 +140,17 @@ def _predict_from_bytes(wav_bytes: bytes) -> dict:
     return _predict(waveform, sr)
 
 
-# ── Endpoints ─────────────────────────────────────────────────────────
-
 @app.get("/")
 @app.get("/health")
 async def health():
     return {
         "status": "ok",
-        "model_loaded": _model is not None,
+        "model_loaded": _model is not None and _model != "dummy",
     }
 
 
 @app.post("/predict")
 async def predict(audio: UploadFile = File(...)):
-    """Predict emotion from an uploaded WAV file."""
     try:
         content = await audio.read()
         if not content:
@@ -146,7 +165,6 @@ async def predict(audio: UploadFile = File(...)):
 
 @app.post("/predict_b64")
 async def predict_b64(data: str = Form(...)):
-    """Predict emotion from base64-encoded WAV bytes."""
     try:
         payload = json.loads(data) if isinstance(data, str) else data
         if isinstance(payload, dict):

requirements.txt

@@ -4,3 +4,6 @@ python-multipart>=0.0.6
 soundfile>=0.12.1
 librosa>=0.10.0
 numpy>=1.24.0
+torch>=2.0.0
+transformers>=4.30.0
+safetensors>=0.3.0