Update app.py

app.py

@@ -108,24 +108,15 @@ def generate_ocr(method, image):
         outputs = model(**inputs)
         logits = outputs.logits  # Get raw logits
 
-    # Print raw logits to debug
+    # Print raw logits for debugging
     print(f"Raw logits: {logits}")
 
-    # Convert logits to probabilities using softmax
-    probs = F.softmax(logits, dim=1)
+    # Compare raw logits instead of using softmax
+    predicted_class = torch.argmax(logits, dim=1).item()
 
-    # Extract probability values
-    not_spam_prob = probs[0, 0].item()
-    spam_prob = probs[0, 1].item()
-
-    # Print probability values for debugging
-    print(f"Not Spam Probability: {not_spam_prob}, Spam Probability: {spam_prob}")
-
-    # Ensure correct label mapping
-    predicted_class = torch.argmax(probs, dim=1).item()  # Get predicted class index
     print(f"Predicted Class Index: {predicted_class}")  # Debugging output
 
-    # Check if the labels are flipped
+    # Ensure correct label mapping
     if predicted_class == 1:
         label = "Spam"
     else: