Irannas/Masked_Email_Classification

1"""Model training and prediction logic for email classification."""
2
3import pickle  # Standard library import first
4import pandas as pd
5from sklearn.feature_extraction.text import TfidfVectorizer
6from sklearn.linear_model import LogisticRegression
7
8
9def train_model():
10    """
11    Train a logistic regression model to classify email types.
12
13    Loads a dataset, preprocesses the text, vectorizes it using TF-IDF, 
14    trains a logistic regression classifier, and saves the model and vectorizer.
15    """
16    df = pd.read_csv("combined_emails_with_natural_pii.csv")
17    df.dropna(subset=["email", "type"], inplace=True)
18
19    email_texts = df["email"]
20    labels = df["type"]
21
22    vectorizer = TfidfVectorizer(
23        stop_words="english", ngram_range=(1, 2), max_df=0.95, min_df=2
24    )
25    vectorized_emails = vectorizer.fit_transform(email_texts)
26
27    model = LogisticRegression(max_iter=1000)
28    model.fit(vectorized_emails, labels)
29
30    with open("classifier_model.pkl", "wb") as f:
31        pickle.dump((model, vectorizer), f)
32
33
34def map_prediction(raw_label):
35    """
36    Map internal model label to a user-friendly category.
37
38    Args:
39        raw_label (str): The raw label predicted by the model.
40
41    Returns:
42        str: Human-readable label for UI display.
43    """
44    label_map = {
45        "Incident": "Technical Support",
46        "Problem": "Billing Issues",
47        "Request": "Account Management",
48        "Change": "Account Management",
49    }
50    return label_map.get(raw_label, raw_label)
51
52
53def predict_category(text):
54    """
55    Predict the category of the given email content.
56
57    Args:
58        text (str): The masked or original email content.
59
60    Returns:
61        str: Predicted email category.
62    """
63    with open("classifier_model.pkl", "rb") as f:
64        model, vectorizer = pickle.load(f)
65
66    vec_text = vectorizer.transform([text])
67    return model.predict(vec_text)[0]
68