Midhran/Email-classification

1import pandas as pd2import joblib3import os4import zipfile5from sklearn.feature_extraction.text import TfidfVectorizer6from sklearn.linear_model import LogisticRegression7from sklearn.model_selection import train_test_split8from sklearn.metrics import classification_report9from sklearn.calibration import CalibratedClassifierCV10from sklearn.preprocessing import LabelEncoder11 12MODEL_PATH = "email_classifier_model.pkl"13VEC_PATH = "tfidf_vectorizer.pkl"14LBL_PATH = "label_encoder.pkl"15ZIP_PATH = "combined_emails_with_natural_pii.zip"16EXTRACT_DIR = "data"17CSV_FILE = "combined_emails_with_natural_pii.csv"18 19def extract_zip(zip_path=ZIP_PATH, extract_to=EXTRACT_DIR):20    if not os.path.exists(extract_to):21        os.makedirs(extract_to)22    with zipfile.ZipFile(zip_path, 'r') as zip_ref:23        zip_ref.extractall(extract_to)24 25def train_classifier(26    save_model=True,27    model_type="logistic"28):29    # Extract the zip file containing the CSV30    extract_zip()31 32    csv_path = os.path.join(EXTRACT_DIR, CSV_FILE)33    df = pd.read_csv(csv_path)34    emails = df["email"]35    labels = df["type"]36 37    X_train, X_test, y_train, y_test = train_test_split(38        emails, labels, test_size=0.2, random_state=4239    )40 41    vectorizer = TfidfVectorizer(max_features=5000)42    X_train_vec = vectorizer.fit_transform(X_train)43    X_test_vec = vectorizer.transform(X_test)44 45    label_encoder = LabelEncoder()46    y_train_enc = label_encoder.fit_transform(y_train)47    y_test_enc = label_encoder.transform(y_test)48 49    model = LogisticRegression(max_iter=300)50    calibrated_model = CalibratedClassifierCV(model)51    calibrated_model.fit(X_train_vec, y_train_enc)52 53    preds = calibrated_model.predict(X_test_vec)54    print("\n--- Classification Report ---")55    print(classification_report(y_test_enc, preds, target_names=label_encoder.classes_))56 57    if save_model:58        joblib.dump(calibrated_model, MODEL_PATH)59        joblib.dump(vectorizer, VEC_PATH)60        joblib.dump(label_encoder, LBL_PATH)61 62    return calibrated_model, vectorizer, label_encoder63 64def load_model():65    if not os.path.exists(MODEL_PATH):66        print("Training model...")67        return train_classifier()68 69    model = joblib.load(MODEL_PATH)70    vectorizer = joblib.load(VEC_PATH)71    label_encoder = joblib.load(LBL_PATH)72    return model, vectorizer, label_encoder73 74def predict_category(email_text):75    model, vectorizer, label_encoder = load_model()76    vec = vectorizer.transform([email_text])77    probs = model.predict_proba(vec)[0]78    pred_idx = probs.argmax()79    label = label_encoder.inverse_transform([pred_idx])[0]80    confidence = float(round(probs[pred_idx], 4))81    return label, confidence82