Complete rewrite: Clean structure with robust visualizations

app.py

@@ -1,14 +1,13 @@
 #!/usr/bin/env python3
 """
-Gradio app for HuggingFace Spaces
-Sentiment analysis with MLflow metrics visualization
+Sentiment Analysis App for HuggingFace Spaces
+Clean, robust implementation with proper visualization
 """
 
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import numpy as np
-import plotly.express as px
 import plotly.graph_objects as go
 import pandas as pd
 from typing import Dict, List, Tuple
@@ -18,50 +17,68 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# ============================================================================
+# SENTIMENT ANALYZER CLASS
+# ============================================================================
+
 class SentimentAnalyzer:
-    """Sentiment analyzer for production"""
+    """Production-ready sentiment analyzer"""
     
-    def __init__(self):
-        # Use the deployed model from HuggingFace Hub
-        self.model_name = "MartinRodrigo/distilbert-sentiment-imdb"
+    def __init__(self, model_name: str = "MartinRodrigo/distilbert-sentiment-imdb"):
+        self.model_name = model_name
         self.tokenizer = None
         self.model = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.load_model()
     
     def load_model(self):
-        """Load the fine-tuned model"""
+        """Load model from HuggingFace Hub"""
         try:
             logger.info(f"Loading model: {self.model_name}")
             self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
             self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
-            logger.info("Model loaded successfully!")
+            self.model.to(self.device)
+            self.model.eval()
+            logger.info(f"✓ Model loaded successfully on {self.device}")
         except Exception as e:
-            logger.error(f"Error loading model: {e}")
-            # Fallback to base model
-            logger.info("Falling back to base model...")
+            logger.error(f"✗ Error loading model: {e}")
+            logger.info("Falling back to base DistilBERT model...")
             self.model_name = "distilbert-base-uncased-finetuned-sst-2-english"
             self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
             self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
+            self.model.to(self.device)
+            self.model.eval()
     
-    def analyze_single(self, text: str) -> Dict:
-        """Analyze sentiment of a single text"""
-        if not text.strip():
+    def predict(self, text: str) -> Dict:
+        """Predict sentiment for a single text"""
+        if not text or not text.strip():
             return {
-                "sentiment": "Please enter some text",
+                "sentiment": "ERROR",
                 "confidence": 0.0,
-                "probabilities": None
+                "probabilities": {"Negative": 0.5, "Positive": 0.5},
+                "error": "Please enter some text"
             }
         
         try:
-            inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+            # Tokenize
+            inputs = self.tokenizer(
+                text, 
+                return_tensors="pt", 
+                truncation=True, 
+                max_length=512,
+                padding=True
+            )
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
             
+            # Predict
             with torch.no_grad():
                 outputs = self.model(**inputs)
-                predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+                probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
             
-            probs = predictions[0].numpy()
-            predicted_class = np.argmax(probs)
-            confidence = float(probs[predicted_class])
+            # Extract results
+            probs_cpu = probs.cpu().numpy()[0]
+            predicted_class = int(np.argmax(probs_cpu))
+            confidence = float(probs_cpu[predicted_class])
             
             sentiment = "POSITIVE" if predicted_class == 1 else "NEGATIVE"
             
@@ -69,270 +86,368 @@ class SentimentAnalyzer:
                 "sentiment": sentiment,
                 "confidence": confidence,
                 "probabilities": {
-                    "Negative": float(probs[0]),
-                    "Positive": float(probs[1])
-                }
+                    "Negative": float(probs_cpu[0]),
+                    "Positive": float(probs_cpu[1])
+                },
+                "error": None
             }
         
         except Exception as e:
-            logger.error(f"Error in analysis: {e}")
+            logger.error(f"Prediction error: {e}")
             return {
-                "sentiment": f"Error: {str(e)}",
+                "sentiment": "ERROR",
                 "confidence": 0.0,
-                "probabilities": None
+                "probabilities": {"Negative": 0.5, "Positive": 0.5},
+                "error": str(e)
             }
     
-    def analyze_batch(self, texts: List[str]) -> List[Dict]:
-        """Analyze multiple texts"""
-        results = []
-        for text in texts:
-            if text.strip():
-                results.append(self.analyze_single(text))
-        return results
+    def predict_batch(self, texts: List[str]) -> List[Dict]:
+        """Predict sentiment for multiple texts"""
+        return [self.predict(text) for text in texts if text.strip()]
 
-# Initialize analyzer
-analyzer = SentimentAnalyzer()
+# ============================================================================
+# VISUALIZATION FUNCTIONS
+# ============================================================================
 
-def analyze_sentiment(text: str) -> Tuple[str, float, go.Figure]:
-    """Main analysis function for Gradio"""
-    result = analyzer.analyze_single(text)
+def create_probability_chart(probabilities: Dict[str, float]) -> go.Figure:
+    """Create a bar chart for sentiment probabilities"""
+    fig = go.Figure()
     
-    # Create figure even if no text to avoid None errors
-    if result["probabilities"]:
-        df = pd.DataFrame([
-            {"Sentiment": "Negative", "Probability": result["probabilities"]["Negative"]},
-            {"Sentiment": "Positive", "Probability": result["probabilities"]["Positive"]}
-        ])
-        
-        fig = px.bar(
-            df, 
-            x="Sentiment", 
-            y="Probability",
-            color="Sentiment",
-            color_discrete_map={"Negative": "#ff4444", "Positive": "#44ff44"},
-            title="Sentiment Probability Distribution",
-            text="Probability"
-        )
-        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%})"
-        
-    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
+    sentiments = list(probabilities.keys())
+    values = list(probabilities.values())
+    colors = ["#ff4444" if s == "Negative" else "#44ff44" for s in sentiments]
+    
+    fig.add_trace(go.Bar(
+        x=sentiments,
+        y=values,
+        marker_color=colors,
+        text=[f"{v:.1%}" for v in values],
+        textposition='outside',
+        textfont=dict(size=14, color='white'),
+        hovertemplate='<b>%{x}</b><br>Probability: %{y:.2%}<extra></extra>'
+    ))
+    
+    fig.update_layout(
+        title={
+            'text': "Sentiment Probability Distribution",
+            'x': 0.5,
+            'xanchor': 'center',
+            'font': {'size': 16, 'color': 'white'}
+        },
+        xaxis_title="Sentiment",
+        yaxis_title="Probability",
+        yaxis_range=[0, 1],
+        height=350,
+        template="plotly_dark",
+        showlegend=False,
+        margin=dict(t=60, b=60, l=60, r=40),
+        paper_bgcolor='rgba(0,0,0,0)',
+        plot_bgcolor='rgba(0,0,0,0)'
+    )
+    
+    return fig
+
+def create_batch_chart(results: List[Dict], texts: List[str]) -> go.Figure:
+    """Create a bar chart for batch analysis results"""
+    fig = go.Figure()
+    
+    text_labels = [f"Text {i+1}" for i in range(len(results))]
+    confidences = [r['confidence'] for r in results]
+    sentiments = [r['sentiment'] for r in results]
+    colors = ["#ff4444" if s == "NEGATIVE" else "#44ff44" for s in sentiments]
+    
+    # Truncate texts for hover
+    hover_texts = [t[:100] + "..." if len(t) > 100 else t for t in texts]
+    
+    fig.add_trace(go.Bar(
+        x=text_labels,
+        y=confidences,
+        marker_color=colors,
+        text=[f"{c:.1%}" for c in confidences],
+        textposition='outside',
+        textfont=dict(size=12, color='white'),
+        customdata=list(zip(sentiments, hover_texts)),
+        hovertemplate='<b>%{x}</b><br>' +
+                      'Sentiment: %{customdata[0]}<br>' +
+                      'Confidence: %{y:.1%}<br>' +
+                      'Text: %{customdata[1]}<extra></extra>'
+    ))
+    
+    fig.update_layout(
+        title={
+            'text': "Batch Analysis Results",
+            'x': 0.5,
+            'xanchor': 'center',
+            'font': {'size': 16, 'color': 'white'}
+        },
+        xaxis_title="Text Number",
+        yaxis_title="Confidence",
+        yaxis_range=[0, 1.1],
+        height=400,
+        template="plotly_dark",
+        showlegend=False,
+        margin=dict(t=60, b=80, l=60, r=40),
+        paper_bgcolor='rgba(0,0,0,0)',
+        plot_bgcolor='rgba(0,0,0,0)'
+    )
+    
+    return fig
+
+# ============================================================================
+# GRADIO INTERFACE FUNCTIONS
+# ============================================================================
+
+def analyze_text(text: str) -> Tuple[str, str, go.Figure]:
+    """Analyze single text and return formatted results"""
+    result = analyzer.predict(text)
+    
+    if result['error']:
+        return (
+            f"⚠️ **Error:** {result['error']}",
+            "0.0",
+            create_probability_chart({"Negative": 0.5, "Positive": 0.5})
         )
-        fig.update_layout(height=300)
-        output_text = result['sentiment']
     
-    return output_text, result['confidence'], fig
+    # Format output
+    emoji = "😊" if result['sentiment'] == "POSITIVE" else "😞"
+    sentiment_text = f"{emoji} **{result['sentiment']}**"
+    confidence_text = f"{result['confidence']:.1%}"
+    
+    # Create chart
+    chart = create_probability_chart(result['probabilities'])
+    
+    return sentiment_text, confidence_text, chart
 
-def analyze_batch_texts(text_input: str) -> Tuple[str, go.Figure]:
-    """Analyze multiple texts separated by newlines"""
-    if not text_input.strip():
+def analyze_batch(text_input: str) -> Tuple[str, go.Figure]:
+    """Analyze multiple texts"""
+    if not text_input or not text_input.strip():
         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(
+            template="plotly_dark",
+            paper_bgcolor='rgba(0,0,0,0)',
+            annotations=[{
+                'text': "Please enter texts to analyze",
+                'xref': 'paper',
+                'yref': 'paper',
+                'x': 0.5,
+                'y': 0.5,
+                'showarrow': False,
+                'font': {'size': 16, 'color': 'gray'}
+            }]
         )
-        empty_fig.update_layout(height=400)
-        return "Please enter some texts (one per line)", empty_fig
+        return "Please enter texts (one per line)", empty_fig
     
+    # Split texts
     texts = [line.strip() for line in text_input.split('\n') if line.strip()]
     
     if not texts:
         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(
+            template="plotly_dark",
+            paper_bgcolor='rgba(0,0,0,0)',
+            annotations=[{
+                'text': "No valid texts found",
+                'xref': 'paper',
+                'yref': 'paper',
+                'x': 0.5,
+                'y': 0.5,
+                'showarrow': False,
+                'font': {'size': 16, 'color': 'gray'}
+            }]
         )
-        empty_fig.update_layout(height=400)
         return "No valid texts found", empty_fig
     
-    results = analyzer.analyze_batch(texts)
+    # Analyze all texts
+    results = analyzer.predict_batch(texts)
     
-    summary_lines = []
-    plot_data = []
-    
-    for i, (text, result) in enumerate(zip(texts, results)):
-        sentiment = result['sentiment']
-        confidence = result['confidence']
-        summary_lines.append(f"{i+1}. **{sentiment}** ({confidence:.1%}) - {text[:50]}{'...' if len(text) > 50 else ''}")
-        
-        plot_data.append({
-            "Text": f"Text {i+1}",
-            "Sentiment": sentiment,
-            "Confidence": confidence
-        })
+    # Create summary
+    summary_lines = ["### 📊 Analysis Results\n"]
+    for i, (text, result) in enumerate(zip(texts, results), 1):
+        emoji = "😊" if result['sentiment'] == "POSITIVE" else "😞"
+        text_preview = text[:60] + "..." if len(text) > 60 else text
+        summary_lines.append(
+            f"{i}. {emoji} **{result['sentiment']}** ({result['confidence']:.1%}) - *{text_preview}*"
+        )
     
     summary = "\n".join(summary_lines)
     
-    # Always create a figure
-    if plot_data:
-        df = pd.DataFrame(plot_data)
-        fig = px.bar(
-            df,
-            x="Text",
-            y="Confidence",
-            color="Sentiment",
-            color_discrete_map={"NEGATIVE": "#ff4444", "POSITIVE": "#44ff44"},
-            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)
+    # Create chart
+    chart = create_batch_chart(results, texts)
     
-    return summary, fig
-    return summary, None
+    return summary, chart
+
+# ============================================================================
+# INITIALIZE ANALYZER
+# ============================================================================
+
+logger.info("Initializing sentiment analyzer...")
+analyzer = SentimentAnalyzer()
+
+# ============================================================================
+# EXAMPLES
+# ============================================================================
 
-# Demo examples
-EXAMPLES = [
-    "🎬 This movie absolutely blew my mind! Best film I've seen this year!",
-    "😞 Worst customer service ever. Total waste of money.",
-    "🚀 Revolutionary AI technology! Incredible understanding of language.",
-    "❌ I regret this purchase deeply. Poor quality materials.",
-    "✈️ Amazing travel experience! The hotel exceeded expectations!",
-    "🎵 Concert was phenomenal! Everything was absolutely perfect!"
+SINGLE_EXAMPLES = [
+    ["This movie is absolutely fantastic! I loved every minute of it."],
+    ["Terrible experience. Waste of time and money."],
+    ["The product is okay, nothing special but it works."],
+    ["Best purchase ever! Highly recommend to everyone!"],
+    ["Disappointed with the quality. Expected much better."]
 ]
 
-BATCH_EXAMPLE = """🛍️ This online store has amazing customer service!
-😡 Terrible experience with their support team.
-⭐ Outstanding quality! Exceeded all my expectations.
-💸 Disappointed with this expensive purchase."""
+BATCH_EXAMPLE = """This movie is absolutely fantastic! I loved every minute of it.
+Terrible experience. Waste of time and money.
+The product is okay, nothing special but it works.
+Best purchase ever! Highly recommend to everyone!"""
+
+# ============================================================================
+# GRADIO UI
+# ============================================================================
 
-# Create Gradio interface
 with gr.Blocks(
-    title="🤖 Transformer Sentiment Analysis",
     theme=gr.themes.Soft(
         primary_hue="blue",
-        secondary_hue="purple",
-        neutral_hue="slate"
-    )
+        secondary_hue="green",
+    ),
+    css="""
+        .gradio-container {
+            max-width: 1200px !important;
+        }
+        #component-0 {
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        }
+    """
 ) as demo:
     
+    # Header
     gr.Markdown("""
-    # 🤖 Transformer Sentiment Analysis
+    # 🎭 Sentiment Analysis with DistilBERT
     
-    Advanced AI-powered sentiment analysis using **DistilBERT** fine-tuned on IMDB dataset.
+    Analyze text sentiment with a fine-tuned transformer model trained on IMDB reviews.
     
-    **Model Performance:**
-    - 🎯 Accuracy: **80%** on test set
-    - 📊 F1 Score: **0.7981**
-    - ⚡ Speed: ~100ms per prediction
-    - 🧠 Parameters: 66M (DistilBERT)
+    **Model:** DistilBERT | **Accuracy:** 80% | **F1:** 0.7981
     """)
     
+    # Main interface
     with gr.Tabs():
+        # TAB 1: Single Analysis
         with gr.TabItem("🔍 Single Analysis"):
-            gr.Markdown("""
-            Analyze individual texts to get sentiment predictions with confidence scores.
-            """)
+            gr.Markdown("### Analyze individual texts")
             
             with gr.Row():
-                with gr.Column(scale=2):
+                with gr.Column(scale=1):
                     single_input = gr.Textbox(
-                        label="💬 Enter your text",
-                        placeholder="Type your text here...",
-                        lines=4
+                        label="Enter your text",
+                        placeholder="Type or paste your text here...",
+                        lines=6,
+                        max_lines=10
+                    )
+                    single_btn = gr.Button(
+                        "🚀 Analyze",
+                        variant="primary",
+                        size="lg"
+                    )
+                    
+                    gr.Examples(
+                        examples=SINGLE_EXAMPLES,
+                        inputs=single_input,
+                        label="Try these examples:"
                     )
-                    single_btn = gr.Button("🚀 Analyze Sentiment", variant="primary", size="lg")
                 
-                with gr.Column(scale=2):
-                    single_output = gr.Markdown(label="📋 Result")
-                    confidence_score = gr.Number(label="🎯 Confidence", precision=3)
-                    probability_plot = gr.Plot(label="📊 Probabilities")
-            
-            gr.Examples(
-                examples=EXAMPLES,
-                inputs=single_input,
-                label="💡 Try these examples:"
-            )
+                with gr.Column(scale=1):
+                    single_sentiment = gr.Markdown(
+                        label="Result",
+                        value="*Results will appear here*"
+                    )
+                    single_confidence = gr.Textbox(
+                        label="Confidence Score",
+                        interactive=False
+                    )
+                    single_plot = gr.Plot(label="Probability Distribution")
         
+        # TAB 2: Batch Processing
         with gr.TabItem("📊 Batch Processing"):
-            gr.Markdown("""
-            Process multiple texts at once. Enter one text per line.
-            """)
+            gr.Markdown("### Process multiple texts at once (one per line)")
             
             with gr.Row():
-                with gr.Column(scale=2):
+                with gr.Column(scale=1):
                     batch_input = gr.Textbox(
-                        label="📝 Multiple texts (one per line)",
+                        label="Enter multiple texts (one per line)",
                         placeholder="Enter texts, one per line...",
-                        lines=8,
+                        lines=10,
                         value=BATCH_EXAMPLE
                     )
-                    batch_btn = gr.Button("🚀 Process Batch", variant="primary", size="lg")
+                    batch_btn = gr.Button(
+                        "🚀 Process Batch",
+                        variant="primary",
+                        size="lg"
+                    )
                 
-                with gr.Column(scale=2):
-                    batch_output = gr.Markdown(label="📈 Results")
-                    batch_plot = gr.Plot(label="📊 Analytics")
+                with gr.Column(scale=1):
+                    batch_results = gr.Markdown(
+                        label="Results Summary",
+                        value="*Results will appear here*"
+                    )
+                    batch_plot = gr.Plot(label="Batch Analytics")
         
+        # TAB 3: About
         with gr.TabItem("ℹ️ About"):
             gr.Markdown("""
             ## About This Model
             
             ### 🏗️ Architecture
-            - **Model:** DistilBERT (Distilled BERT)
+            - **Base Model:** DistilBERT (Distilled BERT)
             - **Parameters:** 66 million
-            - **Training:** Fine-tuned on IMDB dataset
-            - **Accuracy:** 80% on test set
+            - **Training Data:** IMDB Movie Reviews (50k reviews)
+            - **Fine-tuning:** Binary sentiment classification
+            
+            ### 📊 Performance Metrics
+            - **Test Accuracy:** 80.0%
+            - **F1 Score:** 0.7981
+            - **Precision:** High
+            - **Recall:** Balanced
+            
+            ### ⚡ Features
+            - Fast inference (~100ms per prediction)
+            - Batch processing support
+            - Interactive visualizations
+            - Production-ready deployment
             
-            ### ⚡ Performance
-            - **Speed:** ~100ms per prediction
-            - **Batch Processing:** Supported
-            - **Memory:** Optimized for production
+            ### 🔗 Resources
+            - **Model Repository:** [MartinRodrigo/distilbert-sentiment-imdb](https://huggingface.co/MartinRodrigo/distilbert-sentiment-imdb)
+            - **Space:** [transformer-sentiment-analysis](https://huggingface.co/spaces/MartinRodrigo/transformer-sentiment-analysis)
+            - **GitHub:** [ransformer-sentiment-analysis](https://github.com/mrdesautu/ransformer-sentiment-analysis)
             
-            ### 🚀 Tech Stack
+            ### 🛠️ Tech Stack
             - **Framework:** PyTorch + Transformers
-            - **Tracking:** MLflow experiments
             - **UI:** Gradio
-            
-            ### 🔗 Links
-            - **Model:** [MartinRodrigo/distilbert-sentiment-imdb](https://huggingface.co/MartinRodrigo/distilbert-sentiment-imdb)
-            - **GitHub:** [transformer-sentiment-analysis](https://github.com/mrdesautu/ransformer-sentiment-analysis)
+            - **Visualization:** Plotly
+            - **Tracking:** MLflow (local development)
             
             ---
             
-            Built with ❤️ using Transformers, MLflow, and Gradio
+            Built with ❤️ by Martin Rodrigo
             """)
     
-    # Event handlers
+    # Connect event handlers
     single_btn.click(
-        fn=analyze_sentiment,
-        inputs=single_input,
-        outputs=[single_output, confidence_score, probability_plot]
+        fn=analyze_text,
+        inputs=[single_input],
+        outputs=[single_sentiment, single_confidence, single_plot]
     )
     
     batch_btn.click(
-        fn=analyze_batch_texts,
-        inputs=batch_input,
-        outputs=[batch_output, batch_plot]
+        fn=analyze_batch,
+        inputs=[batch_input],
+        outputs=[batch_results, batch_plot]
     )
 
+# ============================================================================
+# LAUNCH
+# ============================================================================
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        share=False,
+        show_error=True
+    )