Update app.py

app.py

@@ -95,7 +95,7 @@ def generate_ocr(method, image):
 
     # Join extracted text into a single string
     text_output = " ".join(extracted_text).strip()
-    
+
     # If no text detected, return early
     if len(text_output) == 0:
         return "No text detected!", "Cannot classify"
@@ -106,11 +106,12 @@ def generate_ocr(method, image):
     # Perform inference
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = F.softmax(outputs.logits, dim=1)  # Convert logits to probabilities
-        spam_prob = probs[0][1].item()  # Probability of Spam
+        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)
 
-    # Adjust classification based on threshold
-    label = "Spam" if spam_prob > 0.6 else "Not Spam"
+    # Map predicted class correctly
+    label = "Spam" if predicted_class == 1 else "Not Spam"
 
     # Save results using external function
     save_results_to_repo(text_output, label)