ENRON email classification [scikit-learn]

Giskard is an open-source framework for testing all ML models, from LLMs to tabular models. Don’t hesitate to give the project a star on GitHub ⭐️ if you find it useful!

In this notebook, you’ll learn how to create comprehensive test suites for your model in a few lines of code, thanks to Giskard’s open-source Python library.

Use-case:

• Multinomial classification of the email’s category.

• Model: LogisticRegression

• Dataset

Outline:

• Detect vulnerabilities automatically with Giskard’s scan

• Automatically generate & curate a comprehensive test suite to test your model beyond accuracy-related metrics

Install dependencies

Make sure to install the giskard

!pip install giskard --upgrade

Import libraries

import glob
import email
from string import punctuation
from collections import defaultdict

import nltk
import pandas as pd
from dateutil import parser
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
from sklearn import model_selection
from sklearn.pipeline import Pipeline
from sklearn.metrics import accuracy_score
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfTransformer

from giskard import Dataset, Model, scan, testing

Define constants

TEXT_COLUMN = "Content"
TARGET_COLUMN = "Target"

COLUMN_TYPES = {TEXT_COLUMN: "text"}
COLUMNS_NAMES = ['Target', 'Subject', 'Content', 'Week_day', 'Year', 'Month', 'Hour', 'Nb_of_forwarded_msg']

RANDOM_STATE = 0

IDX_TO_CAT = {
    1: 'REGULATION',
    2: 'INTERNAL',
    3: 'INFLUENCE',
    4: 'INFLUENCE',
    5: 'INFLUENCE',
    6: 'CALIFORNIA CRISIS',
    7: 'INTERNAL',
    8: 'INTERNAL',
    9: 'INFLUENCE',
    10: 'REGULATION',
    11: 'talking points',
    12: 'meeting minutes',
    13: 'trip reports'}

LABEL_CAT = 3

Dataset preparation

Load and preprocess data

!wget https://bailando.berkeley.edu/enron/enron_with_categories.tar.gz
!tar zxf enron_with_categories.tar.gz
!rm enron_with_categories.tar.gz

nltk.download('punkt')
nltk.download('stopwords')

stoplist = list(set(stopwords.words('english') + list(punctuation)))
stemmer = PorterStemmer()


def get_labels(filename):
    with open(filename + '.cats') as f:
        labels = defaultdict(dict)
        line = f.readline()

        while line:
            line = line.split(',')
            top_cat, sub_cat, freq = int(line[0]), int(line[1]), int(line[2])
            labels[top_cat][sub_cat] = freq
            line = f.readline()

    return dict(labels)


data = pd.DataFrame(columns=COLUMNS_NAMES)

email_files = [f.replace('.cats', '') for f in glob.glob('enron_with_categories/*/*.cats')]
for email_file in email_files:
    values_to_add = {}

    #Target is the sub-category with maximum frequency
    if LABEL_CAT in get_labels(email_file):
        sub_cat_dict = get_labels(email_file)[LABEL_CAT]
        target_int = max(sub_cat_dict, key=sub_cat_dict.get)
        values_to_add[TARGET_COLUMN] = str(IDX_TO_CAT[target_int])

    #Features are metadata from the email object
    filename = email_file + '.txt'
    with open(filename) as f:

        message = email.message_from_string(f.read())

        values_to_add['Subject'] = str(message['Subject'])
        values_to_add[TEXT_COLUMN] = str(message.get_payload())

        date_time_obj = parser.parse(message['Date'])
        values_to_add['Week_day'] = date_time_obj.strftime("%A")
        values_to_add['Year'] = date_time_obj.strftime("%Y")
        values_to_add['Month'] = date_time_obj.strftime("%B")
        values_to_add['Hour'] = int(date_time_obj.strftime("%H"))

        # Count number of forwarded mails
        number_of_messages = 0
        for line in message.get_payload().split('\n'):
            if ('forwarded' in line.lower() or 'original' in line.lower()) and '--' in line:
                number_of_messages += 1
        values_to_add['Nb_of_forwarded_msg'] = number_of_messages

    row_to_add = pd.Series(values_to_add)
    data = pd.concat([data, pd.DataFrame([row_to_add])], ignore_index=True)


data_filtered = data[data[TARGET_COLUMN].notnull()]

# Exclude target category with very few rows.
excluded_category = [IDX_TO_CAT[i] for i in [11, 12, 13]]
data_filtered = data_filtered[data_filtered["Target"].isin(excluded_category) == False]
num_classes = len(data_filtered[TARGET_COLUMN].value_counts())

# Keep only the email column and the target
data_filtered = data_filtered[[TEXT_COLUMN, TARGET_COLUMN]]

Train-test split

Y = data_filtered[TARGET_COLUMN]
X = data_filtered.drop(columns=[TARGET_COLUMN])[list(COLUMN_TYPES.keys())]
X_train, X_test, Y_train, Y_test = model_selection.train_test_split(X, Y, random_state=RANDOM_STATE, stratify=Y)

Wrap dataset with Giskard

To prepare for the vulnerability scan, make sure to wrap your dataset using Giskard’s Dataset class. More details here.

raw_data = pd.concat([X_test, Y_test], axis=1)
giskard_dataset = Dataset(
    df=raw_data,  # A pandas.DataFrame that contains the raw data (before all the pre-processing steps) and the actual ground truth variable (target).
    target="Target",  # Ground truth variable.
    name="Emails of different categories"  # Optional.
)

Model building

Build estimator

text_transformer = Pipeline([
    ('vect', CountVectorizer(stop_words=stoplist)),
    ('tfidf', TfidfTransformer())
])

preprocessor = ColumnTransformer(
    transformers=[
        ('text_Mail', text_transformer, "Content")
    ]
)

clf = Pipeline(steps=[('preprocessor', preprocessor),
                      ('classifier', LogisticRegression())])
clf.fit(X_train, Y_train)

print(f"Train Accuracy score: {clf.score(X_train, Y_train):.2f}")
print(f"Test Accuracy score: {clf.score(X_test, Y_test):.2f}")

Wrap model with Giskard

To prepare for the vulnerability scan, make sure to wrap your model using Giskard’s Model class. You can choose to either wrap the prediction function (preferred option) or the model object. More details here.

giskard_model = Model(
    model=clf,  # A prediction function that encapsulates all the data pre-processing steps and that could be executed with the dataset used by the scan.
    model_type="classification",  # Either regression, classification or text_generation.
    name="Email category classifier",  # Optional.
    classification_labels=clf.classes_.tolist(),  # Their order MUST be identical to the prediction_function's output order.
    feature_names=COLUMN_TYPES.keys(),  # Default: all columns of your dataset.
)

# Validate wrapped model.
y_test_pred_wrapped = giskard_model.predict(giskard_dataset).prediction
wrapped_test_metric = accuracy_score(Y_test, y_test_pred_wrapped)

print(f"Wrapped Test Accuracy score: {wrapped_test_metric:.2f}")

Detect vulnerabilities in your model

Scan your model for vulnerabilities with Giskard

Giskard’s scan allows you to detect vulnerabilities in your model automatically. These include performance biases, unrobustness, data leakage, stochasticity, underconfidence, ethical issues, and more. For detailed information about the scan feature, please refer to our scan documentation.

results = scan(giskard_model, giskard_dataset)

display(results)