Update app.py

app.py

@@ -3,23 +3,20 @@ import spaces
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Initialize GPU
+# Initialize GPU tensor
 zero = torch.Tensor([0]).cuda()
 print(f"Initial device: {zero.device}")
 
 # Load model and tokenizer
 model_name = "tabularisai/robust-sentiment-analysis"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-
-# Move model to GPU
-model = model.to(zero.device)
+model = AutoModelForSequenceClassification.from_pretrained(model_name).cuda()
 
 @spaces.GPU
-def predict_sentiment(text):
+def predict_sentiment(text, show_probabilities=False):
     print(f"Device inside function: {zero.device}")
-    inputs = tokenizer(text.lower(), return_tensors="pt", truncation=True, padding=True, max_length=512)
-    inputs = {k: v.to(zero.device) for k, v in inputs.items()}
+    
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(zero.device)
     
     with torch.no_grad():
         outputs = model(**inputs)
@@ -28,15 +25,37 @@ def predict_sentiment(text):
     predicted_class = torch.argmax(probabilities, dim=-1).item()
     
     sentiment_map = {0: "Very Negative", 1: "Negative", 2: "Neutral", 3: "Positive", 4: "Very Positive"}
-    return sentiment_map[predicted_class]
+    predicted_sentiment = sentiment_map[predicted_class]
+    
+    confidence = probabilities[0][predicted_class].item()
+    
+    result = f"Sentiment: {predicted_sentiment}\nConfidence: {confidence:.2%}\n\n"
+    
+    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"
+    
+    return result
 
-# Gradio interface
+# Create Gradio interface
 demo = gr.Interface(
     fn=predict_sentiment,
-    inputs=gr.Textbox(label="Enter your text here"),
-    outputs=gr.Textbox(label="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!"
+    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()