AkshatJain20092002/Email-Classification-Support-System

1import joblib2import pandas as pd3import logging4from sklearn.feature_extraction.text import TfidfVectorizer5from sklearn.linear_model import LogisticRegression6from sklearn.model_selection import GridSearchCV, train_test_split7from imblearn.pipeline import Pipeline as ImbPipeline8from imblearn.over_sampling import SMOTE9 10 11# Configure logging12logging.basicConfig(13    level=logging.INFO,14    format="%(asctime)s - %(levelname)s - %(message)s",15    datefmt="%Y-%m-%d %H:%M:%S",16)17logger = logging.getLogger(__name__)18 19 20class EmailClassifier:21    """22    Trainer and predictor for email classification.23    """24    def __init__(self):25        self.pipeline = ImbPipeline([26            (27                'tfidf',28                TfidfVectorizer(29                    stop_words='english', max_features=10000,30                ),31            ),32            ('smote', SMOTE(random_state=42)),33            (34                'clf',35                LogisticRegression(36                    solver='saga',37                    class_weight='balanced',38                    max_iter=500,39                    n_jobs=-1,40                ),41            ),42        ])43        self.search = None44 45    def train(self, data_path: str) -> None:46        """47        Train the model and save the best estimator to disk.48        """49        df = pd.read_csv(data_path)50        X = df['email_translated']51        y = df['type']52 53        X_train, _, y_train, _ = train_test_split(54            X, y, test_size=0.2, stratify=y, random_state=42,55        )56 57        param_grid = {'clf__C': [0.1, 1.0, 10.0]}58        self.search = GridSearchCV(59            self.pipeline, param_grid,60            cv=3, scoring='f1_macro', n_jobs=-1, verbose=161        )62        self.search.fit(X_train, y_train)63        joblib.dump(self.search.best_estimator_, 'final_pipeline.pkl')64        logger.info("Model trained and saved as final_pipeline.pkl")65 66    def predict(self, text):67        model = joblib.load('final_pipeline.pkl')68        logger.info("Model loaded from final_pipeline.pkl")69        return model.predict([text])[0]70