Update app.py

app.py

@@ -2,6 +2,7 @@ import evaluate
 from evaluate.utils import launch_gradio_widget
 import gradio as gr
 import torch
+from speechbrain.pretrained.interfaces import foreign_class
 from transformers import AutoModelForSequenceClassification, pipeline, RobertaForSequenceClassification, RobertaTokenizer, AutoTokenizer
 # pull in emotion detection
 # --- Add element for specification
@@ -11,6 +12,15 @@ from transformers import AutoModelForSequenceClassification, pipeline, RobertaFo
 # add logic to initiate mock notificaiton when detected
 # pull in misophonia-specific model
 
+# Building prediction function for gradio
+emotion_dict = {
+    'sad': 'Sad', 
+    'hap': 'Happy',
+    'ang': 'Anger',
+    'neu': 'Neutral'
+}
+
+
 # Create a Gradio interface with audio file and text inputs
 def classify_toxicity(audio_file, text_input, classify_anxiety):
     # Transcribe the audio file using Whisper ASR
@@ -43,8 +53,11 @@ def classify_toxicity(audio_file, text_input, classify_anxiety):
     classification_output = classifier(sequence_to_classify, candidate_labels, multi_label=False)
     print(classification_output)
 
-
-    return toxicity_score, transcribed_text
+    # Emotion classification
+    emotion_classifier = foreign_class(source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP", pymodule_file="custom_interface.py", classname="CustomEncoderWav2vec2Classifier")
+    out_prob, score, index, text_lab = learner.classify_file(audio_file.name)
+ 
+    return toxicity_score, classification_output, emo_dict[text_lab[0]], transcribed_text
     # return f"Toxicity Score ({available_models[selected_model]}): {toxicity_score:.4f}"
  
 with gr.Blocks() as iface: