Update app.py

app.py

@@ -3,29 +3,26 @@ import torch
 import torch.nn as nn
 import joblib
 
-# ---------- Model Definition (Fixed to match .pth file) ----------
+# ---------- Model Definition (Fixed Architecture) ----------
 class SpamClassifier(nn.Module):
     def __init__(self, input_dim):
         super(SpamClassifier, self).__init__()
-        # Based on the error log, the saved model has:
-        # Layer 0: Linear(1000 -> 64)
-        # Layer 2: Linear(64 -> 32)
-        # Layer 4: Linear(32 -> 2)
         self.model = nn.Sequential(
             nn.Linear(input_dim, 64),
             nn.ReLU(),
             nn.Linear(64, 32),
             nn.ReLU(),
-            nn.Linear(32, 2),
-            nn.Softmax(dim=1)
+            # CHANGED: Output layer is 1, not 2
+            nn.Linear(32, 1),
+            # CHANGED: Use Sigmoid for binary output (0 to 1) instead of Softmax
+            nn.Sigmoid()
         )
 
     def forward(self, x):
         return self.model(x)
 
 # ---------- Load Vectorizer ----------
-# Note: Ensure the environment uses scikit-learn version close to 1.2.2 
-# if possible to avoid the warnings in your logs, though it often works anyway.
+# Ensure you have scikit-learn==1.2.2 in requirements.txt if you get warnings
 vectorizer = joblib.load("model/vectorizer.pkl")
 input_dim = len(vectorizer.get_feature_names_out())
 
@@ -38,10 +35,13 @@ model.eval()
 def predict_email(text):
     X = vectorizer.transform([text]).toarray()
     X_tensor = torch.tensor(X, dtype=torch.float32)
+    
     with torch.no_grad():
-        probs = model(X_tensor).numpy()[0]
-    labels = ["Ham", "Spam"]
-    return {labels[i]: float(probs[i]) for i in range(2)}
+        # Get the single probability value (0 = Ham, 1 = Spam)
+        prob_spam = model(X_tensor).item()
+        
+    # Manually calculate Ham probability as the complement of Spam
+    return {"Ham": 1 - prob_spam, "Spam": prob_spam}
 
 # ---------- Gradio Interface ----------
 iface = gr.Interface(