Update app.py

app.py

@@ -8,13 +8,14 @@ import numpy as np
 
 
 
-
+mean, std = -8.278621631819787e-05, 0.08485510250851999
+id2label = {0: 'arousal', 1: 'dominance', 2: 'valence'}
 
 
 
 def classify_audio(audio_file):
     model = AutoModelForAudioClassification.from_pretrained("3loi/SER-Odyssey-Baseline-WavLM-Multi-Attributes", trust_remote_code=True)
-    mean, std = -8.278621631819787e-05, 0.08485510250851999
+    
     
     sr, raw_wav = audio_file
     
@@ -30,8 +31,14 @@ def classify_audio(audio_file):
     wavs = torch.tensor(norm_wav).unsqueeze(0)
     
     pred = model(wavs, mask).detach().numpy()
-    print(str(pred))
-    return str(pred)
+    
+    pred = {}
+    for i, audio_pred in enumerate(pred.numpy()):
+        pred[i] = {}
+        for att_i, att_val in enumerate(audio_pred):
+            pred[i][id2label[att_i]] = att_val
+    
+    return pred
 
 
 def main():