Fixed test

Author: adamspd

spam_detector_ai/classifiers/classifier_map.py

@@ -1,7 +1,7 @@
 from spam_detector_ai.classifiers import ClassifierType, NaiveBayesClassifier, RandomForestSpamClassifier, SVMClassifier
 
 CLASSIFIER_MAP = {
-    ClassifierType.NAIVE_BAYES: NaiveBayesClassifier,
-    ClassifierType.RANDOM_FOREST: RandomForestSpamClassifier,
-    ClassifierType.SVM: SVMClassifier
+    ClassifierType.NAIVE_BAYES.value: NaiveBayesClassifier(),
+    ClassifierType.RANDOM_FOREST.value: RandomForestSpamClassifier(),
+    ClassifierType.SVM.value: SVMClassifier()
 }

spam_detector_ai/prediction/predict.py

@@ -7,6 +7,7 @@
 
 import os
 
+from spam_detector_ai.classifiers import NaiveBayesClassifier, RandomForestSpamClassifier, SVMClassifier
 from spam_detector_ai.classifiers.classifier_map import CLASSIFIER_MAP
 from spam_detector_ai.classifiers.classifier_types import ClassifierType
 from spam_detector_ai.loading_and_processing import Preprocessor
@@ -49,11 +50,18 @@ class SpamDetector:
     """This class is used to detect whether a message is spam or not spam."""
 
     def __init__(self, model_type=ClassifierType.NAIVE_BAYES):
-        classifier_class = CLASSIFIER_MAP.get(model_type)
-        if not classifier_class:
+        classifier_map = {
+            ClassifierType.NAIVE_BAYES.value: NaiveBayesClassifier(),
+            ClassifierType.RANDOM_FOREST.value: RandomForestSpamClassifier(),
+            ClassifierType.SVM.value: SVMClassifier()
+        }
+        classifier = classifier_map.get(model_type.value)
+        print(f"classifier_class: {classifier}")
+        print(f"model_type: {model_type}")
+        if not classifier:
             raise ValueError(f"Invalid model type: {model_type}")
 
-        self.model = classifier_class()
+        self.model = classifier
         model_path, vectoriser_path = get_model_path(model_type)
         self.model.load_model(model_path, vectoriser_path)
         self.processor = Preprocessor()

spam_detector_ai/test_and_tuning/py_test.py

@@ -17,7 +17,7 @@ def test_model():
     current_dir = os.path.dirname(os.path.abspath(__file__))
     base_dir = os.path.dirname(current_dir)
     data_path = os.path.join(base_dir, 'data/spam.csv')
-    initial_trainer = ModelTrainer(data_path=data_path, classifier_type=None, logger=logger)
+    initial_trainer = ModelTrainer(data_path=data_path, logger=logger)
     processed_data = initial_trainer.preprocess_data_()
     _, X_test, _, y_test = train_test_split(processed_data['processed_text'], processed_data['label'],
                                             test_size=0.2, random_state=0)