Update app.py

app.py

@@ -2,60 +2,163 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import math
 
-# Initialize GPU tensor
-zero = torch.Tensor([0]).cuda()
-print(f"Initial device: {zero.device}")
+# Initialize device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
 # Load model and tokenizer
 model_name = "tabularisai/robust-sentiment-analysis"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name).cuda()
+model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
 
-@spaces.GPU
-def predict_sentiment(text, show_probabilities=False):
-    print(f"Device inside function: {zero.device}")
-    
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(zero.device)
+# Define sentiment mapping
+SENTIMENT_MAP = {
+    0: "Very Negative",
+    1: "Negative",
+    2: "Neutral",
+    3: "Positive",
+    4: "Very Positive"
+}
+
+def split_text(text, max_length=512):
+    """
+    Splits the input text into chunks where each chunk has tokens <= max_length.
+    Splitting is done at sentence boundaries to maintain context.
+    """
+    # Simple sentence splitting based on periods. For more accurate splitting, consider using nltk or spacy.
+    sentences = text.split('. ')
+    chunks = []
+    current_chunk = ""
     
-    with torch.no_grad():
-        outputs = model(**inputs)
+    for sentence in sentences:
+        # Add the period back if it was removed during splitting
+        if not sentence.endswith('.'):
+            sentence += '.'
+        
+        # Check if adding the sentence exceeds the max_length
+        encoded = tokenizer.encode(current_chunk + " " + sentence, truncation=False)
+        if len(encoded) > max_length:
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+                current_chunk = sentence
+            else:
+                # Single sentence longer than max_length, force split
+                for i in range(0, len(encoded), max_length):
+                    chunk_tokens = encoded[i:i + max_length]
+                    chunk_text = tokenizer.decode(chunk_tokens, skip_special_tokens=True)
+                    chunks.append(chunk_text.strip())
+                current_chunk = ""
+        else:
+            current_chunk += " " + sentence
     
-    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    predicted_class = torch.argmax(probabilities, dim=-1).item()
+    if current_chunk:
+        chunks.append(current_chunk.strip())
     
-    sentiment_map = {0: "Very Negative", 1: "Negative", 2: "Neutral", 3: "Positive", 4: "Very Positive"}
-    predicted_sentiment = sentiment_map[predicted_class]
+    return chunks
+
+@spaces.GPU
+def analyze_sentiment(text, show_probabilities=False):
+    """
+    Analyzes the sentiment of the input text. If the text exceeds the token limit,
+    it splits the text into chunks and aggregates the results.
+    """
+    try:
+        chunks = split_text(text)
+        all_probabilities = []
+        all_predictions = []
+        detailed_results = ""
+
+        for idx, chunk in enumerate(chunks, 1):
+            inputs = tokenizer(chunk, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
+            
+            with torch.no_grad():
+                outputs = model(**inputs)
+            
+            probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
+            predicted_class = probabilities.argmax()
+            predicted_sentiment = SENTIMENT_MAP[predicted_class]
+            confidence = probabilities[predicted_class]
+            
+            all_probabilities.append(probabilities)
+            all_predictions.append(predicted_class)
+            
+            if show_probabilities:
+                detailed_results += f"**Chunk {idx}: {predicted_sentiment} ({confidence:.2%})**\n"
+                for cls, prob in zip(SENTIMENT_MAP.values(), probabilities):
+                    detailed_results += f"{cls}: {prob:.2%}\n"
+                detailed_results += "\n"
+            else:
+                detailed_results += f"**Chunk {idx}: {predicted_sentiment} ({confidence:.2%})**\n"
+        
+        # Aggregate results by averaging probabilities
+        avg_probabilities = sum(all_probabilities) / len(all_probabilities)
+        final_class = avg_probabilities.argmax()
+        final_sentiment = SENTIMENT_MAP[final_class]
+        final_confidence = avg_probabilities[final_class]
+        
+        result = f"**Overall Sentiment: {final_sentiment}**\nConfidence: {final_confidence:.2%}\n\n"
+        
+        if show_probabilities:
+            result += "### Detailed Analysis:\n" + detailed_results
+        else:
+            result += "### Detailed Analysis:\n" + detailed_results
+        
+        return result
+    except Exception as e:
+        return f"An error occurred during sentiment analysis: {str(e)}"
+
+# Create Gradio interface using Blocks for better layout control
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎭 Sentiment Analysis Wizard")
+    gr.Markdown(
+        """
+        Discover the emotional tone behind any text with our advanced AI model! This app uses a state-of-the-art language model to analyze the sentiment of your text, classifying it into one of five categories: **Very Negative**, **Negative**, **Neutral**, **Positive**, or **Very Positive**.
+        """
+    )
     
-    confidence = probabilities[0][predicted_class].item()
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(
+                lines=10,
+                placeholder="Enter text for sentiment analysis...",
+                label="Input Text"
+            )
+            show_probs = gr.Checkbox(
+                label="Show probabilities for each class",
+                value=False
+            )
+            analyze_button = gr.Button("Analyze Sentiment")
+        with gr.Column():
+            output = gr.Markdown(label="Result")
     
-    result = f"Sentiment: {predicted_sentiment}\nConfidence: {confidence:.2%}\n\n"
+    with gr.Accordion("Examples", open=False):
+        examples = [
+            ["I absolutely loved this movie! The acting was superb and the plot was engaging.", True],
+            ["The service at this restaurant was terrible. I'll never go back.", False],
+            ["The product works as expected. Nothing special, but it gets the job done.", True],
+            ["I'm somewhat disappointed with my purchase. It's not as good as I hoped.", False],
+            ["This book changed my life! I couldn't put it down and learned so much.", True]
+        ]
+        gr.Examples(
+            examples=examples,
+            inputs=[input_text, show_probs],
+            label="Predefined Examples"
+        )
     
-    if show_probabilities:
-        result += "Probabilities for each class:\n"
-        for i, (sentiment, prob) in enumerate(zip(sentiment_map.values(), probabilities[0])):
-            result += f"{sentiment}: {prob.item():.2%}\n"
+    analyze_button.click(
+        fn=analyze_sentiment,
+        inputs=[input_text, show_probs],
+        outputs=output
+    )
     
-    return result
-
-# Create Gradio interface
-demo = gr.Interface(
-    fn=predict_sentiment,
-    inputs=[
-        gr.Textbox(lines=5, label="Enter text for sentiment analysis"),
-        gr.Checkbox(label="Show probabilities for each class")
-    ],
-    outputs=gr.Textbox(label="Result"),
-    title="🎭 Sentiment Analysis Wizard",
-    description="Discover the emotional tone behind any text with our advanced AI model! This app uses a state-of-the-art language model to analyze the sentiment of your text, classifying it into one of five categories: Very Negative, Negative, Neutral, Positive, or Very Positive.",
-    examples=[
-        ["I absolutely loved this movie! The acting was superb and the plot was engaging.", True],
-        ["The service at this restaurant was terrible. I'll never go back.", False],
-        ["The product works as expected. Nothing special, but it gets the job done.", True],
-        ["I'm somewhat disappointed with my purchase. It's not as good as I hoped.", False],
-        ["This book changed my life! I couldn't put it down and learned so much.", True]
-    ],
-    theme=gr.themes.Soft()
-)
-
-demo.launch()
+    gr.Markdown(
+        """
+        ---
+        **Developed with ❤️ using Gradio and Transformers by Hugging Face**
+        """
+    )
+
+# Launch the interface
+demo.launch()