Update app.py

app.py

@@ -1,52 +1,51 @@
-import gradio as gr
-import torch
-import torch.nn as nn
-import pickle
-
-# ---------- Model Definition ----------
-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)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.relu(x)
-        x = self.fc2(x)
-        x = self.softmax(x)
-        return x
-
-# ---------- Load Vectorizer ----------
-with open("model/vectorizer.pkl", "rb") as f:
-    vectorizer = pickle.load(f)
-
-input_dim = len(vectorizer.get_feature_names_out())
-
-# ---------- Load Model ----------
-model = SpamClassifier(input_dim)
-model.load_state_dict(torch.load("model/email_spam_classifier.pth", map_location=torch.device("cpu")))
-model.eval()
-
-# ---------- Prediction Function ----------
-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)}
-
-# ---------- Gradio Interface ----------
-iface = gr.Interface(
-    fn=predict_email,
-    inputs=gr.Textbox(lines=5, placeholder="Paste your email here..."),
-    outputs=gr.Label(num_top_classes=2),
-    title="Email Spam Classifier",
-    description="Classify emails as Spam or Ham using a PyTorch model."
-)
-
-if __name__ == "__main__":
-    iface.launch()
+import gradio as gr
+import torch
+import torch.nn as nn
+import joblib
+
+# ---------- Model Definition ----------
+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)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.softmax(x)
+        return x
+
+# ---------- Load Vectorizer ----------
+vectorizer = joblib.load("model/vectorizer.pkl")  # Use joblib for TF-IDF
+
+input_dim = len(vectorizer.get_feature_names_out())
+
+# ---------- Load Model ----------
+model = SpamClassifier(input_dim)
+model.load_state_dict(torch.load("model/email_spam_classifier.pth", map_location=torch.device("cpu")))
+model.eval()
+
+# ---------- Prediction Function ----------
+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)}
+
+# ---------- Gradio Interface ----------
+iface = gr.Interface(
+    fn=predict_email,
+    inputs=gr.Textbox(lines=5, placeholder="Paste your email here..."),
+    outputs=gr.Label(num_top_classes=2),
+    title="Email Spam Classifier",
+    description="Classify emails as Spam or Ham using a PyTorch model."
+)
+
+if __name__ == "__main__":
+    iface.launch()