Optimized functions & changed from pickle to joblib

Author: adamspd

.github/workflows/tests.yml

@@ -38,3 +38,5 @@ jobs:
     - name: Test with pytest
       run: |
         pytest
+        python -m unittest spam_detector_ai.tests.test_dataloader
+        python -m unittest spam_detector_ai.tests.test_preprocessor

README.md

@@ -139,12 +139,12 @@ The test results are shown below:
 |                  | Predicted: Ham       | Predicted: Spam     |
 |------------------|----------------------|---------------------|
 | **Actual: Ham**  | 2080 (True Negative) | 25 (False Positive) |
-| **Actual: Spam** | 42 (False Negative)  | 812 (True Positive) |
+| **Actual: Spam** | 41 (False Negative)  | 813 (True Positive) |
 
 - True Negative (TN): 2080 messages were correctly identified as ham (non-spam).
 - False Positive (FP): 25 ham messages were incorrectly identified as spam.
-- False Negative (FN): 42 spam messages were incorrectly identified as ham.
-- True Positive (TP): 812 messages were correctly identified as spam.
+- False Negative (FN): 41 spam messages were incorrectly identified as ham.
+- True Positive (TP): 813 messages were correctly identified as spam.
 
 ##### Performance Metrics:
 
@@ -212,7 +212,7 @@ The project contains 3 pre-trained models that can be used directly if you want
 If you don't want to use the package, you can use the API that I have deployed
 [here](https://spam-detection-api.adamspierredavid.com/).
 
-The API is built with Django and the following is an example of how I use it in a personal project:
+The API is built with Django, and the following is an example of how I use it in a personal project:
 
 ![Screenshot](./screenshots/spam-detection-api-example.png)
 

requirements.txt

@@ -5,4 +5,5 @@ nltk~=3.8.1
 setuptools==69.0.3
 pytest==7.4.4
 requests~=2.31.0
-imblearn~=0.0
+imblearn~=0.0
+joblib~=1.3.2

spam_detector_ai/classifiers/base_classifier.py

@@ -2,15 +2,28 @@
 
 from abc import ABC, abstractmethod
 
+from joblib import dump, load
+
 
 class BaseClassifier(ABC):
+    VECTORIZER_PARAMS = {
+        'max_features': 1500,
+        'min_df': 5,
+        'max_df': 0.7
+    }
+
     def __init__(self):
-        pass
+        self.classifier = None
+        self.vectoriser = None
 
     @abstractmethod
     def train(self, X_train, y_train):
         pass
 
-    @abstractmethod
     def save_model(self, model_path, vectoriser_path):
-        pass
+        dump(self.classifier, model_path)
+        dump(self.vectoriser, vectoriser_path)
+
+    def load_model(self, model_path, vectoriser_path):
+        self.classifier = load(model_path)
+        self.vectoriser = load(vectoriser_path)

spam_detector_ai/classifiers/classifier_map.py

@@ -0,0 +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
+}

spam_detector_ai/classifiers/naive_bayes_classifier.py

@@ -1,7 +1,5 @@
 # spam_detector_ai/classifiers/naive_bayes_classifier.py
 
-import pickle
-
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.naive_bayes import MultinomialNB
 
@@ -11,16 +9,9 @@
 class NaiveBayesClassifier(BaseClassifier):
     def __init__(self):
         super().__init__()
-        self.classifier = None
-        self.vectoriser = CountVectorizer(max_features=1500, min_df=5, max_df=0.7)
+        self.vectoriser = CountVectorizer(**BaseClassifier.VECTORIZER_PARAMS)
 
     def train(self, X_train, y_train):
         X_train_vectorized = self.vectoriser.fit_transform(X_train).toarray()
         self.classifier = MultinomialNB()
         self.classifier.fit(X_train_vectorized, y_train)
-
-    def save_model(self, model_path, vectoriser_path):
-        with open(model_path, 'wb') as file:
-            pickle.dump(self.classifier, file)
-        with open(vectoriser_path, 'wb') as file:
-            pickle.dump(self.vectoriser, file)

spam_detector_ai/classifiers/random_forest_classifier.py

@@ -1,27 +1,20 @@
 # spam_detector_ai/classifiers/random_forest_classifier.py
 
-import pickle
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.ensemble import RandomForestClassifier
 from imblearn.over_sampling import SMOTE
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.feature_extraction.text import TfidfVectorizer
+
 from .base_classifier import BaseClassifier
 
 
 class RandomForestSpamClassifier(BaseClassifier):
     def __init__(self):
         super().__init__()
-        self.classifier = None
-        self.vectoriser = TfidfVectorizer(max_features=1500, min_df=5, max_df=0.7)
+        self.vectoriser = TfidfVectorizer(**BaseClassifier.VECTORIZER_PARAMS)
         self.smote = SMOTE(random_state=42)
 
     def train(self, X_train, y_train):
         X_train_vectorized = self.vectoriser.fit_transform(X_train)
         X_train_res, y_train_res = self.smote.fit_resample(X_train_vectorized, y_train)
         self.classifier = RandomForestClassifier(n_estimators=100, random_state=0)
         self.classifier.fit(X_train_res, y_train_res)
-
-    def save_model(self, model_path, vectoriser_path):
-        with open(model_path, 'wb') as file:
-            pickle.dump(self.classifier, file)
-        with open(vectoriser_path, 'wb') as file:
-            pickle.dump(self.vectoriser, file)

spam_detector_ai/classifiers/svm_classifier.py

@@ -1,24 +1,17 @@
 # spam_detector_ai/classifiers/svm_classifier.py
 
-import pickle
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.svm import SVC
+
 from .base_classifier import BaseClassifier
 
 
 class SVMClassifier(BaseClassifier):
     def __init__(self):
         super().__init__()
-        self.classifier = None
-        self.vectoriser = TfidfVectorizer(max_features=1500, min_df=5, max_df=0.7)
+        self.vectoriser = TfidfVectorizer(**BaseClassifier.VECTORIZER_PARAMS)
 
     def train(self, X_train, y_train):
         X_train_vectorized = self.vectoriser.fit_transform(X_train)
-        self.classifier = SVC(C=100, gamma=1, kernel='rbf')
+        self.classifier = SVC(C=10, gamma=1, kernel='rbf')
         self.classifier.fit(X_train_vectorized, y_train)
-
-    def save_model(self, model_path, vectoriser_path):
-        with open(model_path, 'wb') as file:
-            pickle.dump(self.classifier, file)
-        with open(vectoriser_path, 'wb') as file:
-            pickle.dump(self.vectoriser, file)

spam_detector_ai/loading_and_processing/__init__.py

@@ -1,4 +1,4 @@
 # spam_detector_ai/loading_and_processing/__init__.py
 
 from .data_loader import DataLoader
-from .preprocessor import Preprocessor
+from .preprocessor import Preprocessor

spam_detector_ai/loading_and_processing/data_loader.py

@@ -5,7 +5,17 @@
 
 class DataLoader:
     def __init__(self, data_path):
-        self.data = pd.read_csv(data_path)
+        if not data_path.endswith('.csv'):
+            raise ValueError("Only CSV files are supported")
+        try:
+            self.data = pd.read_csv(data_path)
+        except FileNotFoundError:
+            raise FileNotFoundError(f"The file at {data_path} was not found.")
+        except Exception as e:
+            raise Exception(f"An error occurred while loading the file: {e}")
 
     def get_data(self):
+        """
+        Return the loaded data.
+        """
         return self.data