manjusha-r/email-classification-api

1import pandas as pd2from sklearn.feature_extraction.text import TfidfVectorizer3from sklearn.naive_bayes import MultinomialNB4from sklearn.pipeline import Pipeline5from sklearn.model_selection import train_test_split6import joblib7import os8 9# Path to dataset (update as needed)10data_path = "data/emails.csv"11 12if not os.path.exists(data_path):13    raise FileNotFoundError(f"{data_path} not found. Please add your dataset to the 'data' folder.")14 15# Load dataset16df = pd.read_csv(data_path)17 18# Check columns19if 'email' not in df.columns or 'category' not in df.columns:20    raise ValueError("Dataset must have 'email' and 'category' columns.")21 22X = df['email']23y = df['category']24 25# Split the data26X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)27 28# Pipeline29model = Pipeline([30    ('tfidf', TfidfVectorizer()),31    ('clf', MultinomialNB())32])33 34# Train model35model.fit(X_train, y_train)36 37# Save model38os.makedirs("model", exist_ok=True)39joblib.dump(model, "model/email_classifier.pkl")40 41print("✅ Model trained and saved in model/email_classifier.pkl")42