Update app.py

app.py

@@ -4,8 +4,12 @@ import evaluate
 from evaluate.utils import launch_gradio_widget
 import gradio as gr
 import torch
+import classify
+from whisper.tokenizer import get_tokenizer
 from speechbrain.pretrained.interfaces import foreign_class
-from transformers import AutoModelForSequenceClassification, pipeline, RobertaForSequenceClassification, RobertaTokenizer, AutoTokenizer
+from transformers import AutoModelForSequenceClassification, pipeline, WhisperTokenizer, RobertaForSequenceClassification, RobertaTokenizer, AutoTokenizer
+
+
 # pull in emotion detection
 # --- Add element for specification
 # pull in text classification
@@ -14,6 +18,8 @@ from transformers import AutoModelForSequenceClassification, pipeline, RobertaFo
 # add logic to initiate mock notificaiton when detected
 # pull in misophonia-specific model
 
+model_cache = {}
+
 # Building prediction function for gradio
 emo_dict = {
     'sad': 'Sad', 
@@ -76,8 +82,38 @@ def classify_toxicity(audio_file, text_input, classify_anxiety):
         return toxicity_score, classification_output, emo_dict[text_lab[0]], transcribed_text
         # return f"Toxicity Score ({available_models[selected_model]}): {toxicity_score:.4f}"
     else: 
+        # model = model_cache[model_name]
+        class_names = classify_anxiety.split(",")
+        # tokenizer = get_tokenizer(multilingual=".en" not in model_name)
+        tokenizer= WhisperTokenizer.from_pretrained("openai/whisper-large")
+        model = "large"
+        
+        if model_name not in model_cache:
+            model = whisper.load_model(model_name)
+            model_cache[model_name] = model
+        else:
+            # model = model_cache[model_name]
+            class_names = classify_anxiety.split(",")
+    
+        internal_lm_average_logprobs = classify.calculate_internal_lm_average_logprobs(
+            model=model,
+            class_names=class_names,
+            class_names=classify_anxiety
+            tokenizer=tokenizer,
+        )
+        audio_features = classify.calculate_audio_features(audio_path, model)
+        average_logprobs = classify.calculate_average_logprobs(
+            model=model,
+            audio_features=audio_features,
+            class_names=class_names,
+            tokenizer=tokenizer,
+        )
+        average_logprobs -= internal_lm_average_logprobs
+        scores = average_logprobs.softmax(-1).tolist()
+        return {class_name: score for class_name, score in zip(class_names, scores)}
+
         return classify_anxiety
- 
+     
 with gr.Blocks() as iface:
     with gr.Column():
         classify = gr.Radio(["racism", "LGBTQ+ hate", "sexually explicit", "misophonia"])