Refactor sentiment prediction function to improve clarity and confidence extraction

app.py

@@ -11,11 +11,26 @@ tokenizer = BertTokenizer.from_pretrained("./imdb_bert_model")
 
 # Prediction function
 def predict_sentiment(text):
+    """
+    Predicts the sentiment of the given text using the fine-tuned BERT model.
+
+    Args:
+        text (str): The input movie review text.
+
+    Returns:
+        str: The predicted sentiment with confidence.
+    """
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     outputs = model(**inputs)
     logits = outputs.logits
-    confidence = torch.max(logits, dim=1).item()
-    prediction = torch.argmax(logits, dim=1).item()
+
+    # Extract the maximum value (confidence) and its index (prediction)
+    confidence, prediction = torch.max(logits, dim=1)
+    confidence = confidence.item()  # Convert tensor to Python float
+    prediction = prediction.item()  # Convert tensor to Python int
+
+    # confidence = torch.max(logits, dim=1).item()
+    # prediction = torch.argmax(logits, dim=1).item()
     sentiment = "Positive 😊" if prediction == 1 else "Negative 😠"
     return f"{sentiment} with confidence {confidence * 100:.2f}% confidence"