training/trainer.py

"""Training loop shared by all four classifiers."""

import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.pipeline import Pipeline

from training.evaluation import evaluate_predictions, select_best_model


def build_text_pipeline(model):
    """Wrap ``model`` behind a ``TfidfVectorizer``."""
    return Pipeline([('vect', TfidfVectorizer()),
                     ('clf', model)])


def train_and_evaluate_models(models, X_train, X_test, y_train, y_test):
    """Train every model in ``models`` and compare them on the test set."""
    results = []

    for model_name, model in models.items():
        # Create a text classification pipeline with TF-IDF vectorizer and the specified model
        text_clf = build_text_pipeline(model)

        # Train the model
        text_clf.fit(X_train, y_train)

        # Make predictions on the test set
        y_pred = text_clf.predict(X_test)

        # Evaluate the model
        metrics = evaluate_predictions(y_test, y_pred)

        # Store the results in a dictionary
        result_dict = {
            'Model': model_name,
            'Trained Model': text_clf,  # Store the trained model
            **metrics,
        }

        # Append the results to the list
        results.append(result_dict)

    # Convert the list of dictionaries to a DataFrame
    results_df = pd.DataFrame(results)

    # Find the best model based on the highest F1 Score
    best_model_name, best_model = select_best_model(results_df)

    return best_model_name, best_model, results_df