Update app.py

app.py

@@ -10,38 +10,23 @@ text_model = AutoModelForSequenceClassification.from_pretrained("tae898/emoberta
 text_model.gradient_checkpointing_enable()
 
 
-def gradio_combined(audio_file, text):
-    # Case 1 — Audio provided
-    if audio_file is not None:
-        waveform, sr = torchaudio.load(audio_file)
-        preds = speech_classifier(waveform.squeeze().numpy(), sampling_rate=sr, top_k=3)
+def predict_emotion(audio, text):
+    results = {}
 
-        return {
-            "Detected Emotion": preds[0]["label"],
-            "Top Predictions": {p["label"]: round(p["score"], 3) for p in preds},
-            "Source": "Audio"
-        }
+    if audio is not None:
+        waveform, sr = torchaudio.load(audio)
+        preds = speech_classifier(waveform.squeeze().numpy(), sampling_rate=sr, top_k=3)
+        results["audio_emotion"] = preds[0]["label"]
 
-    # Case 2 — Text provided
-    if text.strip() != "":
-        inputs = text_tokenizer(text, return_tensors="pt", truncation=True)
+    if text is not None and text.strip() != "":
+        inputs = text_tokenizer(text, return_tensors="pt")
         with torch.no_grad():
             outputs = text_model(**inputs)
+        emotion = text_model.config.id2label[torch.argmax(outputs.logits)]
+        results["text_emotion"] = emotion
 
-        probs = torch.nn.functional.softmax(outputs.logits, dim=1)
-        label_id = torch.argmax(probs).item()
-
-        return {
-            "Detected Emotion": text_model.config.id2label[label_id],
-            "Top Predictions": {
-                text_model.config.id2label[i]: round(p, 3)
-                for i, p in enumerate(probs[0].tolist())
-            },
-            "Source": "Text"
-        }
-
-    return {"Error": "Please provide audio or text input."}
-
+    return results
+    
 # Building the UI
 gradio_ui = gr.Interface(
     fn=gradio_combined,