Fix: Improve probability plots with labels and better error handling

app.py

@@ -93,10 +93,11 @@ class SentimentAnalyzer:
 # Initialize analyzer
 analyzer = SentimentAnalyzer()
 
-def analyze_sentiment(text: str) -> Tuple[str, float, dict]:
+def analyze_sentiment(text: str) -> Tuple[str, float, go.Figure]:
     """Main analysis function for Gradio"""
     result = analyzer.analyze_single(text)
     
+    # Create figure even if no text to avoid None errors
     if result["probabilities"]:
         df = pd.DataFrame([
             {"Sentiment": "Negative", "Probability": result["probabilities"]["Negative"]},
@@ -109,27 +110,56 @@ def analyze_sentiment(text: str) -> Tuple[str, float, dict]:
             y="Probability",
             color="Sentiment",
             color_discrete_map={"Negative": "#ff4444", "Positive": "#44ff44"},
-            title="Sentiment Probability Distribution"
+            title="Sentiment Probability Distribution",
+            text="Probability"
         )
-        fig.update_layout(showlegend=False, height=300)
+        fig.update_traces(texttemplate='%{text:.2%}', textposition='outside')
+        fig.update_layout(
+            showlegend=False, 
+            height=300,
+            yaxis_range=[0, 1],
+            yaxis_title="Probability",
+            xaxis_title="Sentiment"
+        )
+        
+        output_text = f"**{result['sentiment']}** (Confidence: {result['confidence']:.1%})"
         
-        return (
-            f"**{result['sentiment']}** (Confidence: {result['confidence']:.1%})",
-            result['confidence'],
-            fig
+    else:
+        # Create empty figure for error cases
+        fig = go.Figure()
+        fig.add_annotation(
+            text="No data to display",
+            xref="paper", yref="paper",
+            x=0.5, y=0.5, showarrow=False
         )
+        fig.update_layout(height=300)
+        output_text = result['sentiment']
     
-    return result['sentiment'], result['confidence'], None
+    return output_text, result['confidence'], fig
 
-def analyze_batch_texts(text_input: str) -> Tuple[str, dict]:
+def analyze_batch_texts(text_input: str) -> Tuple[str, go.Figure]:
     """Analyze multiple texts separated by newlines"""
     if not text_input.strip():
-        return "Please enter some texts (one per line)", None
+        empty_fig = go.Figure()
+        empty_fig.add_annotation(
+            text="Please enter texts to analyze",
+            xref="paper", yref="paper",
+            x=0.5, y=0.5, showarrow=False
+        )
+        empty_fig.update_layout(height=400)
+        return "Please enter some texts (one per line)", empty_fig
     
     texts = [line.strip() for line in text_input.split('\n') if line.strip()]
     
     if not texts:
-        return "No valid texts found", None
+        empty_fig = go.Figure()
+        empty_fig.add_annotation(
+            text="No valid texts found",
+            xref="paper", yref="paper",
+            x=0.5, y=0.5, showarrow=False
+        )
+        empty_fig.update_layout(height=400)
+        return "No valid texts found", empty_fig
     
     results = analyzer.analyze_batch(texts)
     
@@ -149,6 +179,7 @@ def analyze_batch_texts(text_input: str) -> Tuple[str, dict]:
     
     summary = "\n".join(summary_lines)
     
+    # Always create a figure
     if plot_data:
         df = pd.DataFrame(plot_data)
         fig = px.bar(
@@ -157,12 +188,26 @@ def analyze_batch_texts(text_input: str) -> Tuple[str, dict]:
             y="Confidence",
             color="Sentiment",
             color_discrete_map={"NEGATIVE": "#ff4444", "POSITIVE": "#44ff44"},
-            title="Batch Analysis Results"
+            title="Batch Analysis Results",
+            text="Confidence"
+        )
+        fig.update_traces(texttemplate='%{text:.1%}', textposition='outside')
+        fig.update_layout(
+            height=400,
+            yaxis_range=[0, 1.1],
+            yaxis_title="Confidence",
+            xaxis_title="Text Number"
+        )
+    else:
+        fig = go.Figure()
+        fig.add_annotation(
+            text="No results to display",
+            xref="paper", yref="paper",
+            x=0.5, y=0.5, showarrow=False
         )
         fig.update_layout(height=400)
-        
-        return summary, fig
     
+    return summary, fig
     return summary, None
 
 # Demo examples