Update app.py

app.py

@@ -1,12 +1,14 @@
 import gradio as gr
 import pandas as pd
 import plotly.express as px
+import plotly.graph_objects as go
 import shutil
 import os
 import torch
 from huggingface_hub import hf_hub_download
 from importlib import import_module
 
+# Load inference.py and model
 repo_id = "logasanjeev/goemotions-bert"
 local_file = hf_hub_download(repo_id=repo_id, filename="inference.py")
 print("Downloaded inference.py successfully!")
@@ -24,15 +26,21 @@ _, _ = 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):
+# Prediction function with export capability
+def predict_emotions_with_details(text, confidence_threshold=0.0, chart_type="bar"):
+    if not text.strip():
+        return "Please enter some text.", "", "", None, None
+    
     predictions_str, processed_text = predict_emotions(text)
     
+    # Parse predictions
     predictions = []
     if predictions_str != "No emotions predicted.":
         for line in predictions_str.split("\n"):
             emotion, confidence = line.split(": ")
             predictions.append((emotion, float(confidence)))
     
+    # Get raw logits for all emotions
     encodings = inference_module.TOKENIZER(
         processed_text,
         padding='max_length',
@@ -47,11 +55,13 @@ def predict_emotions_with_details(text, confidence_threshold=0.0):
         outputs = inference_module.MODEL(input_ids, attention_mask=attention_mask)
         logits = torch.sigmoid(outputs.logits).cpu().numpy()[0]
     
+    # All emotions for top 5
     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])
     
+    # Filter predictions based on threshold
     filtered_predictions = []
     for emotion, confidence in predictions:
         thresh = default_thresholds[emotion_labels.index(emotion)]
@@ -64,67 +74,107 @@ def predict_emotions_with_details(text, confidence_threshold=0.0):
     else:
         thresholded_output = "\n".join([f"{emotion}: {confidence:.4f}" for emotion, confidence in filtered_predictions])
     
+    # Create visualization
+    fig = None
+    df_export = None
     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
+        df_export = df.copy()
+        
+        if chart_type == "bar":
+            fig = px.bar(
+                df,
+                x="Emotion",
+                y="Confidence",
+                color="Emotion",
+                text="Confidence",
+                title="Emotion Confidence Levels (Above Threshold)",
+                height=400,
+                color_discrete_sequence=px.colors.qualitative.Plotly
+            )
+            fig.update_traces(texttemplate='%{text:.2f}', textposition='auto')
+            fig.update_layout(showlegend=False, margin=dict(t=40, b=40), xaxis_title="", yaxis_title="Confidence")
+        else:  # pie chart
+            fig = px.pie(
+                df,
+                names="Emotion",
+                values="Confidence",
+                title="Emotion Confidence Distribution (Above Threshold)",
+                height=400,
+                color_discrete_sequence=px.colors.qualitative.Plotly
+            )
+            fig.update_traces(textinfo='percent+label', pull=[0.1] + [0] * (len(df) - 1))
+            fig.update_layout(margin=dict(t=40, b=40))
     
-    return processed_text, thresholded_output, top_5_output, fig
+    return processed_text, thresholded_output, top_5_output, fig, df_export
 
+# Custom CSS for enhanced styling
 custom_css = """
 body {
-    font-family: 'Segoe UI', Arial, sans-serif;
+    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
+    background-color: #f5f7fa;
 }
 .gr-panel {
-    border-radius: 12px;
-    box-shadow: 0 4px 12px rgba(0,0,0,0.1);
-    background: linear-gradient(145deg, #ffffff, #f0f4f8);
+    border-radius: 16px;
+    box-shadow: 0 6px 20px rgba(0,0,0,0.08);
+    background: white;
+    padding: 20px;
+    margin-bottom: 20px;
 }
 .gr-button {
     border-radius: 8px;
-    background: #007bff;
+    padding: 12px 24px;
+    font-weight: 600;
+    transition: all 0.3s ease;
+}
+.gr-button-primary {
+    background: #4a90e2;
     color: white;
-    padding: 10px 20px;
-    transition: background 0.3s;
 }
-.gr-button:hover {
-    background: #0056b3;
+.gr-button-primary:hover {
+    background: #357abd;
+}
+.gr-button-secondary {
+    background: #e4e7eb;
+    color: #333;
+}
+.gr-button-secondary:hover {
+    background: #d1d5db;
 }
 #title {
-    font-size: 2.5em;
+    font-size: 2.8em;
+    font-weight: 700;
     color: #1a3c6e;
     text-align: center;
-    margin-bottom: 20px;
+    margin-bottom: 10px;
 }
 #description {
-    font-size: 1.1em;
-    color: #333;
+    font-size: 1.2em;
+    color: #555;
     text-align: center;
-    max-width: 700px;
-    margin: 0 auto;
+    max-width: 800px;
+    margin: 0 auto 30px auto;
 }
 #theme-toggle {
-    position: absolute;
+    position: fixed;
     top: 20px;
     right: 20px;
+    background: none;
+    border: none;
+    font-size: 1.5em;
+    cursor: pointer;
+    transition: transform 0.3s;
+}
+#theme-toggle:hover {
+    transform: scale(1.2);
 }
 .dark-mode {
-    background: #1a1a1a;
+    background: #1e2a44;
     color: #e0e0e0;
 }
 .dark-mode .gr-panel {
-    background: linear-gradient(145deg, #2a2a2a, #3a3a3a);
+    background: #2a3a5a;
+    box-shadow: 0 6px 20px rgba(0,0,0,0.2);
 }
 .dark-mode #title {
     color: #66b3ff;
@@ -132,80 +182,176 @@ body {
 .dark-mode #description {
     color: #b0b0b0;
 }
+.dark-mode .gr-button-secondary {
+    background: #3a4a6a;
+    color: #e0e0e0;
+}
+.dark-mode .gr-button-secondary:hover {
+    background: #4a5a7a;
+}
+#loading {
+    font-style: italic;
+    color: #888;
+    text-align: center;
+}
+#examples-title {
+    font-size: 1.5em;
+    font-weight: 600;
+    color: #1a3c6e;
+    margin-bottom: 10px;
+}
+.dark-mode #examples-title {
+    color: #66b3ff;
+}
+footer {
+    text-align: center;
+    margin-top: 40px;
+    padding: 20px;
+    font-size: 0.9em;
+    color: #666;
+}
+footer a {
+    color: #4a90e2;
+    text-decoration: none;
+}
+footer a:hover {
+    text-decoration: underline;
+}
+.dark-mode footer {
+    color: #b0b0b0;
+}
 """
 
+# JavaScript for theme toggle and loading spinner
 theme_js = """
 function toggleTheme() {
     document.body.classList.toggle('dark-mode');
+    const toggleBtn = document.getElementById('theme-toggle');
+    toggleBtn.innerHTML = document.body.classList.contains('dark-mode') ? '☀️' : '🌙';
+}
+function showLoading() {
+    document.getElementById('loading').style.display = 'block';
+}
+function hideLoading() {
+    document.getElementById('loading').style.display = 'none';
 }
 """
 
 # Gradio Blocks UI
 with gr.Blocks(css=custom_css) as demo:
+    # Theme toggle button
+    gr.HTML(
+        """
+        <button id='theme-toggle' onclick='toggleTheme()'>🌙</button>
+        <script>{}</script>
+        """.format(theme_js)
+    )
+    
     # Header
     gr.Markdown("<div id='title'>GoEmotions BERT Classifier</div>", elem_id="title")
     gr.Markdown(
         """
         <div id='description'>
-        Predict emotions from text using a fine-tuned BERT-base model. 
-        Explore 28 emotions with optimized thresholds (Micro F1: 0.6006). 
-        View preprocessed text, top 5 emotions, and thresholded predictions!
+        Predict emotions from text using a fine-tuned BERT-base model on the GoEmotions dataset. 
+        Detect 28 emotions with optimized thresholds (Micro F1: 0.6006). 
+        View preprocessed text, top 5 emotions, and thresholded predictions with interactive visualizations!
         </div>
         """,
         elem_id="description"
     )
     
-    # Theme toggle button
+    # Main content
     with gr.Row():
-        gr.HTML(
-            """
-            <button id='theme-toggle' onclick='toggleTheme()'>Toggle Dark Mode</button>
-            <script>{}</script>
-            """.format(theme_js)
-        )
+        with gr.Column(scale=1):
+            # Input Section
+            with gr.Group():
+                gr.Markdown("### Input Text")
+                text_input = gr.Textbox(
+                    label="Enter Your Text",
+                    placeholder="Type something like 'I’m just chilling today'...",
+                    lines=3,
+                    show_label=False
+                )
+                confidence_slider = gr.Slider(
+                    minimum=0.0,
+                    maximum=0.9,
+                    value=0.0,
+                    step=0.05,
+                    label="Minimum Confidence Threshold",
+                    info="Filter predictions below this confidence level (default thresholds still apply)"
+                )
+                chart_type = gr.Radio(
+                    choices=["bar", "pie"],
+                    value="bar",
+                    label="Chart Type",
+                    info="Choose how to visualize the emotion confidences"
+                )
+                with gr.Row():
+                    submit_btn = gr.Button("Predict Emotions", variant="primary")
+                    reset_btn = gr.Button("Reset", variant="secondary")
     
-    # Main input and output
+    # Loading indicator
+    gr.HTML("<div id='loading' style='display:none;'>Predicting emotions, please wait...</div>")
+    
+    # Output Section
     with gr.Row():
         with gr.Column(scale=1):
-            text_input = gr.Textbox(
-                label="Enter Your Text",
-                placeholder="Type something like 'I’m just chilling today'...",
-                lines=3
-            )
-            confidence_slider = gr.Slider(
-                minimum=0.0,
-                maximum=0.9,
-                value=0.0,
-                step=0.05,
-                label="Minimum Confidence Threshold",
-                info="Adjust to filter low-confidence predictions"
-            )
-            submit_btn = gr.Button("Predict Emotions", variant="primary")
-        
-        with gr.Column(scale=1):
-            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)")
+            with gr.Group():
+                gr.Markdown("### Results")
+                processed_text_output = gr.Textbox(label="Preprocessed Text", lines=2, interactive=False)
+                thresholded_output = gr.Textbox(label="Predicted Emotions (Above Threshold)", lines=5, interactive=False)
+                top_5_output = gr.Textbox(label="Top 5 Emotions (Regardless of Threshold)", lines=5, interactive=False)
+                output_plot = gr.Plot(label="Emotion Confidence Visualization (Above Threshold)")
+    
+    # Export predictions
+    export_btn = gr.File(label="Download Predictions as CSV", visible=False)
     
     # Example carousel
-    examples = gr.Examples(
-        examples=[
-            "I’m just chilling today.",
-            "Thank you for saving my life!",
-            "I’m nervous about my exam tomorrow.",
-            "I love my new puppy so much!",
-            "I’m so relieved the storm passed."
-        ],
-        inputs=text_input,
-        label="Try These Examples"
+    with gr.Group():
+        gr.Markdown("<div id='examples-title'>Example Texts</div>", elem_id="examples-title")
+        examples = gr.Examples(
+            examples=[
+                ["I’m just chilling today.", "Neutral Example"],
+                ["Thank you for saving my life!", "Gratitude Example"],
+                ["I’m nervous about my exam tomorrow.", "Nervousness Example"],
+                ["I love my new puppy so much!", "Love Example"],
+                ["I’m so relieved the storm passed.", "Relief Example"]
+            ],
+            inputs=[text_input],
+            label="",
+            examples_per_page=3
+        )
+    
+    # Footer
+    gr.HTML(
+        """
+        <footer>
+            Built with ❤️ by logasanjeev | 
+            <a href="https://huggingface.co/logasanjeev/goemotions-bert">Model Card</a> | 
+            <a href="https://www.kaggle.com/code/ravindranlogasanjeev/evaluation-logasanjeev-goemotions-bert/notebook">Kaggle Notebook</a> | 
+            <a href="https://github.com/logasanjeev">GitHub</a>
+        </footer>
+        """
     )
     
-    # Bind prediction
+    # Bind predictions with loading spinner
     submit_btn.click(
         fn=predict_emotions_with_details,
-        inputs=[text_input, confidence_slider],
-        outputs=[processed_text_output, thresholded_output, top_5_output, output_plot]
+        inputs=[text_input, confidence_slider, chart_type],
+        outputs=[processed_text_output, thresholded_output, top_5_output, output_plot, export_btn],
+        _js="showLoading(); return [arguments[0], arguments[1], arguments[2]]"
+    ).then(
+        fn=None,
+        inputs=None,
+        outputs=None,
+        _js="hideLoading"
+    )
+    
+    # Reset functionality
+    reset_btn.click(
+        fn=lambda: ("", "", "", None, None),
+        inputs=[],
+        outputs=[text_input, processed_text_output, thresholded_output, top_5_output, output_plot, export_btn]
     )
 
 # Launch