Update app.py

app.py

@@ -108,10 +108,19 @@ def generate_ocr(method, image):
         outputs = model(**inputs)
         logits = outputs.logits  # Get raw logits
         probs = F.softmax(logits, dim=1)  # Convert logits to probabilities
-        predicted_class = torch.argmax(probs, dim=1).item()  # Get the predicted class (0 or 1)
 
-    # Map predicted class correctly
-    label = "Spam" if predicted_class == 1 else "Not Spam"
+    # Print probabilities for debugging
+    print(f"Probabilities: {probs}")
+
+    # Extract probability values
+    not_spam_prob = probs[0, 0].item()
+    spam_prob = probs[0, 1].item()
+
+    # Ensure classification is based on correct threshold
+    if spam_prob > not_spam_prob:
+        label = "Spam"
+    else:
+        label = "Not Spam"
 
     # Save results using external function
     save_results_to_repo(text_output, label)