Building an AI For a Spam Classifier (With Python)

Classical AI and neural nets built in python can be used for multiple infosec tasks for defensive measures. The cool part: They don't need a GPU or a super powerful setup. This time: A Spam Classifier.

The steps for building a spam classifier in python are:

1. Get a labeled dataset of spam and non-spam (aka ham) messages
2. Standardize/Normalize the messages
3. Tokenize the messages
4. Feature Extraction
5. Train a model
6. Evaluate the model

A Naive-Bayes approach for this example will be the perfect fit. (See here as to why). Essentially we will calculate the probability of something being spam given the words that it contains.

For the naive bayes to work we have to build a "histogram" first, so we need data about spam and ham messages.

Such a dataset could look like this for example:

-------------------- HEAD --------------------
  label                                            message
0   ham  Go until jurong point, crazy.. Available only ...
1   ham                      Ok lar... Joking wif u oni...
2  spam  Free entry in 2 a wkly comp to win FA Cup fina...
3   ham  U dun say so early hor... U c already then say...
4   ham  Nah I don't think he goes to usf, he lives aro...

Source

01 Standardizing the data

To make the model efficient we have to reduce and simplify the dataset. The goal is to make it possible for the model to understand "the context" of a message.

First, we deduplicate the data. Using python and pandas we can load a dataframe object with our data and then run

# Remove duplicates if any
df = df.drop_duplicates()

to remove the duplicates. Then we lowercase the messages with

df["message"] = df["message"].str.lower()

For the model to understand the context, it doesn't matter if the message contains "FREE BEER" or "free beer". Both cases imply that the message talks about the free beer. (Even tho for us as humans a message all in caps seems more sus).

With

df["message"] = df["message"].apply(lambda x: re.sub(r"[^a-z\s$!]", "", x))

We remove everything other the $ sign and the ! (and of course all of the letters). We do this so the model can focus on the words, to get the context/intent of the message.

And while "You won 1'000'000$" looks more like spam than "You have an outstanding bill of 24.99$", we don't wanna include every possible price tag, since that would make the model to granular. If we keep all the numbers it may be able to detect 1'000'000$ as spam, but not know what to make of 999'999$.

02 Tokenizing the data

After we prepped our text more or less, we can tokenize it (while skipping the stop words) using

import nltk

# Download the necessary NLTK data files
nltk.download("punkt")
nltk.download("punkt_tab")
nltk.download("stopwords")

from nltk.tokenize import word_tokenize

# Split each message into individual tokens
df["message"] = df["message"].apply(word_tokenize)

from nltk.corpus import stopwords

# Define a set of English stop words and remove them from the tokens
stop_words = set(stopwords.words("english"))
df["message"] = df["message"].apply(lambda x: [word for word in x if word not in stop_words])

We can do this since words like and, in, the, etc... don't really add to the context of the message for a spam detection.

Next we can also apply stemming. Stemming is basically dumbing down a text. It applies a normalization of words to their base: i.e. running -> run. This makes the model understand the underlying concepts of the sentences, makes them more comparable and is more efficient.

from nltk.stem import PorterStemmer

# Stem each token to reduce words to their base form
stemmer = PorterStemmer()
df["message"] = df["message"].apply(lambda x: [stemmer.stem(word) for word in x])

After we tokenized the messages and applied more simplifications we can join them back together as strings, since most AI/Data science libraries work best with strings that they tokenize themselves.

df["message"] = df["message"].apply(lambda x: " ".join(x))

03 Feature Extraction

The next step is to transform the text into numerical representations. Models can't work with text, so we need to convert the messages into numeric representations.

Bag-of-words

A common way to represent text numerically is by creating a vector of term counts. Each element in the vector then corresponds to a word/term in the vocabulary of all the messages.

We can use unigrams for single words and also introduce bigrams to include sense of order or meaning. (For example "free ipad" can help us more in distinguishing spam than just "free" and "ipad")

Using scikit-learn we can apply tokenization, build a vocabulary and map each document to a numeric vector.

We use the following key params:

• min_df=1 Word must appear in at least one document (can be set higher to exclude rare terms)
• max_df=0.9 Very common words (appearing in 90% of the documents) are excluded
• ngram_range=(1,2) telling it to include unigrams and bigrams

from sklearn.feature_extraction.text import CountVectorizer

# Initialize CountVectorizer with bigrams, min_df, and max_df to focus on relevant terms
vectorizer = CountVectorizer(min_df=1, max_df=0.9, ngram_range=(1, 2))

# Fit and transform the message column
X = vectorizer.fit_transform(df["message"])

# Labels (target variable)
y = df["label"].apply(lambda x: 1 if x == "spam" else 0)  # Converting labels to 1 and 0

In the end we are gonna have a vector for each document that tells us how often certain words are in that message if that message appears in at least one document.

For example the vector [1,2,3,2,7,0,1] could mean that the document with the nr. 1 contains 2 occurences of "free", 3 occurences of "ipad", 2 occurences of "you won", 7 occurences of "you" and 0 of "money". The last 1 means it is spam. (Simplistic example)

The model internally defines in which order the words (and which words) will end up in those vectors, and each vector contains the same amount of words in the same order.

04 Training the Model

Now we finished preparing the data for a model to work with it: We can start training a model using the Naive Bayes Classification.

We Streamline it into a pipeline. This allows us to apply the steps of vectorization and data transformation before applying the Bayes. It also allows us to do hyperparamenter tuning.

In this example we use GridsearchCV which will try the different values for alpha (part of the naive bayes formula).

from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline

# Build the pipeline by combining vectorization and classification
pipeline = Pipeline([
    ("vectorizer", vectorizer),
    ("classifier", MultinomialNB())
])

# Define the parameter grid for hyperparameter tuning
param_grid = {
    "classifier__alpha": [0.01, 0.1, 0.15, 0.2, 0.25, 0.5, 0.75, 1.0]
}

# Perform the grid search with 5-fold cross-validation and the F1-score as metric
grid_search = GridSearchCV(
    pipeline,
    param_grid,
    cv=5,
    scoring="f1"
)

# Fit the grid search on the full dataset
grid_search.fit(df["message"], y)

# Extract the best model identified by the grid search
best_model = grid_search.best_estimator_

05 Evaluation

The last step is to evaluate the model. For this we take a test-sample of data (or create one) and evaluate it with our model.

These messages need to go through all the same modifications as the training data -> otherwise the model can't compare it.

# Messages
new_messages = [
    "Congratulations! You've won a $1000 Walmart gift card. Go to http://bit.ly/1234 to claim now.",
    "Hey, are we still meeting up for lunch today?",
    "Urgent! Your account has been compromised. Verify your details here: www.fakebank.com/verify",
    "Reminder: Your appointment is scheduled for tomorrow at 10am.",
    "FREE entry in a weekly competition to win an iPad. Just text WIN to 80085 now!",
]

import numpy as np
import re

# Preprocess function that mirrors the training-time preprocessing
def preprocess_message(message):
    message = message.lower()
    message = re.sub(r"[^a-z\s$!]", "", message)
    tokens = word_tokenize(message)
    tokens = [word for word in tokens if word not in stop_words]
    tokens = [stemmer.stem(word) for word in tokens]
    return " ".join(tokens)

# Preprocess and vectorize messages
processed_messages = [preprocess_message(msg) for msg in new_messages]

# Transform preprocessed messages into feature vectors
X_new = best_model.named_steps["vectorizer"].transform(processed_messages)

# Predict with the trained classifier
predictions = best_model.named_steps["classifier"].predict(X_new)
prediction_probabilities = best_model.named_steps["classifier"].predict_proba(X_new)

# Display predictions and probabilities for each evaluated message
for i, msg in enumerate(new_messages):
    prediction = "Spam" if predictions[i] == 1 else "Not-Spam"
    spam_probability = prediction_probabilities[i][1]  # Probability of being spam
    ham_probability = prediction_probabilities[i][0]   # Probability of being not spam
    
    print(f"Message: {msg}")
    print(f"Prediction: {prediction}")
    print(f"Spam Probability: {spam_probability:.2f}")
    print(f"Not-Spam Probability: {ham_probability:.2f}")
    print("-" * 50)

This will show for each message how likely it is spam. For example:

Message: Congratulations! You've won a $1000 Walmart gift card. Go to http://bit.ly/1234 to claim now.
Prediction: Spam
Spam Probability: 1.00
Not-Spam Probability: 0.00

Saving the model

If you're happy with the result/model you can save it using joblib.

import joblib

# Save the trained model to a file for future use
model_filename = 'spam_detection_model.joblib'
joblib.dump(best_model, model_filename)

print(f"Model saved to {model_filename}")

*This study note is based on the AI Red Teamer job path modules/materials I'm studying for my Certified Offensive AI Expert certification I'm studying for from HTB.*