Update app.py

app.py

@@ -10,6 +10,7 @@ import easyocr
 import keras_ocr
 from paddleocr import PaddleOCR
 from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
+import torch.nn.functional as F  # Added for softmax
 
 # Paths
 MODEL_PATH = "./distilbert_spam_model"
@@ -19,7 +20,7 @@ RESULTS_CSV = "ocr_results.csv"
 # Ensure model exists
 if not os.path.exists(os.path.join(MODEL_PATH, "pytorch_model.bin")):
     print(f"⚠️ Model not found in {MODEL_PATH}. Downloading from Hugging Face Hub...")
-    model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")
+    model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
     model.save_pretrained(MODEL_PATH)
     tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
     tokenizer.save_pretrained(MODEL_PATH)
@@ -63,11 +64,15 @@ def generate_ocr(method, img):
         text_output = ocr_with_keras(img)
 
     # Classify Text as Spam or Not Spam
-    inputs = tokenizer(text_output, return_tensors="pt", truncation=True, padding=True)
+    inputs = tokenizer(text_output, return_tensors="pt", truncation=True, padding=True, max_length=512)
+
     with torch.no_grad():
         outputs = model(**inputs)
-    prediction = torch.argmax(outputs.logits, dim=1).item()
-    label = "Spam" if prediction == 1 else "Not Spam"
+        probs = F.softmax(outputs.logits, dim=1)  # Convert logits to probabilities
+        prediction = torch.argmax(probs, dim=1).item()
+
+    label_map = {0: "Not Spam", 1: "Spam"}
+    label = label_map[prediction]
 
     # Save results
     save_results(text_output, label)