Fixed word2vec vectorizer for SVM model

advaithsrao · Nov 8, 2023 · 3e28e6e · 3e28e6e
1 parent b6b06c2
commit 3e28e6e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/detector/modeler.py b/detector/modeler.py
@@ -339,8 +339,8 @@ def train(
             label = label.tolist()
 
         # Vectorize the input texts
-        X = self.vectorizer.fit_transform(body)
-        y = np.array(label)
+        X = [self.vectorizer.fit_transform(doc) for doc in body]
+        y = label
 
         # Train the SVM model
         self.model.fit(X, y)
@@ -361,7 +361,7 @@ def predict(
             body = body.tolist()
 
         # Vectorize the input texts
-        X = self.vectorizer.transform(body)
+        X = [self.vectorizer.fit_transform(doc) for doc in body]
 
         # Make predictions using the trained SVM model
         predictions = self.model.predict(X)