Update app.py

app.py

@@ -1,76 +1,86 @@
 import gradio as gr
-import torch
-import numpy as np
-from transformers import BertForSequenceClassification, BertTokenizer
-import requests
-import json
-import plotly.express as px
 import pandas as pd
+import plotly.express as px
+import shutil
+import os
+from huggingface_hub import hf_hub_download
+from importlib import import_module
 
-# Load model and tokenizer from Hugging Face Hub
 repo_id = "logasanjeev/goemotions-bert"
-model = BertForSequenceClassification.from_pretrained(repo_id)
-tokenizer = BertTokenizer.from_pretrained(repo_id)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
-if torch.cuda.device_count() > 1:
-    model = nn.DataParallel(model)
-model.eval()
+local_file = hf_hub_download(repo_id=repo_id, filename="inference.py")
+print("Downloaded inference.py successfully!")
+
+current_dir = os.getcwd()
+destination = os.path.join(current_dir, "inference.py")
+shutil.copy(local_file, destination)
+print("Copied inference.py to current directory!")
 
-# Load optimized thresholds from Hugging Face Hub
-thresholds_url = f"https://huggingface.co/{repo_id}/raw/main/thresholds.json"
-response = requests.get(thresholds_url)
-thresholds_data = json.loads(response.text)
-emotion_labels = thresholds_data["emotion_labels"]
-default_thresholds = thresholds_data["thresholds"]
+inference_module = import_module("inference")
+predict_emotions = inference_module.predict_emotions
+print("Imported predict_emotions successfully!")
 
-# Prediction function
-def predict_emotions(text, confidence_threshold=0.0):
-    encodings = tokenizer(
-        text,
+_, _ = predict_emotions("dummy text")
+emotion_labels = inference_module.EMOTION_LABELS
+default_thresholds = inference_module.THRESHOLDS
+
+def predict_emotions_with_details(text, confidence_threshold=0.0):
+    predictions_str, processed_text = predict_emotions(text)
+    
+    predictions = []
+    if predictions_str != "No emotions predicted.":
+        for line in predictions_str.split("\n"):
+            emotion, confidence = line.split(": ")
+            predictions.append((emotion, float(confidence)))
+    
+    encodings = inference_module.TOKENIZER(
+        processed_text,
         padding='max_length',
         truncation=True,
         max_length=128,
         return_tensors='pt'
     )
-    input_ids = encodings['input_ids'].to(device)
-    attention_mask = encodings['attention_mask'].to(device)
+    input_ids = encodings['input_ids'].to(inference_module.DEVICE)
+    attention_mask = encodings['attention_mask'].to(inference_module.DEVICE)
     
     with torch.no_grad():
-        outputs = model(input_ids, attention_mask=attention_mask)
+        outputs = inference_module.MODEL(input_ids, attention_mask=attention_mask)
         logits = torch.sigmoid(outputs.logits).cpu().numpy()[0]
     
-    # Apply thresholds with user-defined confidence boost
-    predictions = []
-    for i, (logit, thresh) in enumerate(zip(logits, default_thresholds)):
-        adjusted_thresh = max(thresh, confidence_threshold)
-        if logit >= adjusted_thresh:
-            predictions.append((emotion_labels[i], logit))
+    all_emotions = [(emotion_labels[i], round(logit, 4)) for i, logit in enumerate(logits)]
+    all_emotions.sort(key=lambda x: x[1], reverse=True)
+    top_5_emotions = all_emotions[:5]
+    top_5_output = "\n".join([f"{emotion}: {confidence:.4f}" for emotion, confidence in top_5_emotions])
     
-    predictions.sort(key=lambda x: x[1], reverse=True)
-    if not predictions:
-        return "No emotions predicted above thresholds.", None
+    filtered_predictions = []
+    for emotion, confidence in predictions:
+        thresh = default_thresholds[emotion_labels.index(emotion)]
+        adjusted_thresh = max(thresh, confidence_threshold)
+        if confidence >= adjusted_thresh:
+            filtered_predictions.append((emotion, confidence))
     
-    # Format output
-    text_output = "\n".join([f"{emotion}: {confidence:.4f}" for emotion, confidence in predictions])
+    if not filtered_predictions:
+        thresholded_output = "No emotions predicted above thresholds."
+    else:
+        thresholded_output = "\n".join([f"{emotion}: {confidence:.4f}" for emotion, confidence in filtered_predictions])
     
-    # Create bar chart
-    df = pd.DataFrame(predictions, columns=["Emotion", "Confidence"])
-    fig = px.bar(
-        df,
-        x="Emotion",
-        y="Confidence",
-        color="Emotion",
-        text="Confidence",
-        title="Emotion Confidence Levels",
-        height=400
-    )
-    fig.update_traces(texttemplate='%{text:.2f}', textposition='auto')
-    fig.update_layout(showlegend=False, margin=dict(t=40, b=40))
+    if filtered_predictions:
+        df = pd.DataFrame(filtered_predictions, columns=["Emotion", "Confidence"])
+        fig = px.bar(
+            df,
+            x="Emotion",
+            y="Confidence",
+            color="Emotion",
+            text="Confidence",
+            title="Emotion Confidence Levels (Above Threshold)",
+            height=400
+        )
+        fig.update_traces(texttemplate='%{text:.2f}', textposition='auto')
+        fig.update_layout(showlegend=False, margin=dict(t=40, b=40))
+    else:
+        fig = None
     
-    return text_output, fig
+    return processed_text, thresholded_output, top_5_output, fig
 
-# Custom CSS for modern UI
 custom_css = """
 body {
     font-family: 'Segoe UI', Arial, sans-serif;
@@ -123,7 +133,6 @@ body {
 }
 """
 
-# JavaScript for theme toggle
 theme_js = """
 function toggleTheme() {
     document.body.classList.toggle('dark-mode');
@@ -138,8 +147,8 @@ with gr.Blocks(css=custom_css) as demo:
         """
         <div id='description'>
         Predict emotions from text using a fine-tuned BERT-base model. 
-        Explore 28 emotions with optimized thresholds (Micro F1: 0.6025). 
-        Try examples or enter your own text!
+        Explore 28 emotions with optimized thresholds (Micro F1: 0.6006). 
+        View preprocessed text, top 5 emotions, and thresholded predictions!
         </div>
         """,
         elem_id="description"
@@ -173,8 +182,10 @@ with gr.Blocks(css=custom_css) as demo:
             submit_btn = gr.Button("Predict Emotions", variant="primary")
         
         with gr.Column(scale=1):
-            output_text = gr.Textbox(label="Predicted Emotions", lines=5)
-            output_plot = gr.Plot(label="Emotion Confidence Chart")
+            processed_text_output = gr.Textbox(label="Preprocessed Text", lines=2)
+            thresholded_output = gr.Textbox(label="Predicted Emotions (Above Threshold)", lines=5)
+            top_5_output = gr.Textbox(label="Top 5 Emotions (Regardless of Threshold)", lines=5)
+            output_plot = gr.Plot(label="Emotion Confidence Chart (Above Threshold)")
     
     # Example carousel
     examples = gr.Examples(
@@ -191,9 +202,9 @@ with gr.Blocks(css=custom_css) as demo:
     
     # Bind prediction
     submit_btn.click(
-        fn=predict_emotions,
+        fn=predict_emotions_with_details,
         inputs=[text_input, confidence_slider],
-        outputs=[output_text, output_plot]
+        outputs=[processed_text_output, thresholded_output, top_5_output, output_plot]
     )
 
 # Launch