Update app.py

Update change HTML to Images.

app.py

@@ -4,6 +4,7 @@ import torch
 import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
+from io import BytesIO
 
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment")
@@ -11,16 +12,8 @@ model = AutoModelForSequenceClassification.from_pretrained("nlptown/bert-base-mu
 
 # Define prediction function
 def predict(texts):
-    processed_texts = []
-    for text in texts:
-        if isinstance(text, list):
-            processed_text = tokenizer.convert_tokens_to_string(text)
-        else:
-            processed_text = text
-        processed_texts.append(processed_text)
-    
     inputs = tokenizer(
-        processed_texts, 
+        texts, 
         return_tensors="pt", 
         padding=True, 
         truncation=True, 
@@ -48,25 +41,28 @@ def analyze_text(text):
     # Generate SHAP explanations
     shap_values = explainer([text])
     
-    # Create HTML visualizations for all classes
-    html_plots = []
-    for i in range(shap_values.shape[-1]):
-        # Generate SHAP plot HTML
-        plot_html = shap.plots.text(shap_values[0, :, i], display=False)
+    # Create matplotlib figures for each class
+    images = []
+    for i in range(5):
+        plt.figure(figsize=(10, 3))
+        shap.plots.bar(shap_values[0, :, i], show=False)
+        plt.title(f"Feature importance for {output_names_list[i]}")
+        plt.tight_layout()
         
-        # Create HTML visualizations for all classes
-        html_plots = []
-        for i in range(shap_values.shape[-1]):
-            # Create SHAP text plot and convert to HTML
-            plot_html = shap.plots.text(shap_values[0, :, i], display=False)
-            html_plots.append(plot_html)
-        # Format confidence scores
-        confidence_scores = {model.config.id2label[i]: float(probabilities[i]) 
-                        for i in range(len(probabilities))}
+        # Save plot to in-memory buffer
+        buf = BytesIO()
+        plt.savefig(buf, format="png", bbox_inches="tight")
+        plt.close()
+        buf.seek(0)
+        images.append(buf)
+    
+    # Format confidence scores
+    confidence_scores = {model.config.id2label[i]: float(probabilities[i]) 
+                         for i in range(len(probabilities))}
     
-    return (predicted_label, confidence_scores, *html_plots)
+    return (predicted_label, confidence_scores, *images)
 
-# Create Gradio interface with HTML components
+# Create Gradio interface with image components
 with gr.Blocks() as demo:
     gr.Markdown("## 🔍 BERT Sentiment Analysis with SHAP Explanations")
     
@@ -84,7 +80,7 @@ with gr.Blocks() as demo:
         gr.Markdown("""
         ### SHAP Explanations
         Below you can see how each word contributes to different sentiment scores (1-5 stars).  
-        Red text increases the score, blue decreases it.
+        Positive values increase the score, negative values decrease it.
         """)
     
     # Individual Explanation Rows
@@ -92,8 +88,9 @@ with gr.Blocks() as demo:
     for i in range(5):
         with gr.Row():
             plot_components.append(
-                gr.HTML(
+                gr.Image(
                     label=f"Explanation for {model.config.id2label[i]}",
+                    type="pil",
                     elem_classes=f"shap-plot-{i+1}"
                 )
             )
@@ -115,4 +112,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share = True, ssr_mode = False, )
+    demo.launch(share=True)