Update app.py

app.py

@@ -3,25 +3,22 @@ import torch
 import torch.nn as nn
 import joblib
 
-# ---------- Model Definition ----------
+# ---------- Model Definition (matches saved model) ----------
 class SpamClassifier(nn.Module):
     def __init__(self, input_dim):
         super(SpamClassifier, self).__init__()
-        self.fc1 = nn.Linear(input_dim, 128)
-        self.relu = nn.ReLU()
-        self.fc2 = nn.Linear(128, 2)
-        self.softmax = nn.Softmax(dim=1)
+        self.model = nn.Sequential(
+            nn.Linear(input_dim, 128),
+            nn.ReLU(),
+            nn.Linear(128, 2),
+            nn.Softmax(dim=1)
+        )
 
     def forward(self, x):
-        x = self.fc1(x)
-        x = self.relu(x)
-        x = self.fc2(x)
-        x = self.softmax(x)
-        return x
+        return self.model(x)
 
 # ---------- Load Vectorizer ----------
-vectorizer = joblib.load("model/vectorizer.pkl")  # Use joblib for TF-IDF
-
+vectorizer = joblib.load("model/vectorizer.pkl")
 input_dim = len(vectorizer.get_feature_names_out())
 
 # ---------- Load Model ----------