training/evaluation.py

"""Metrics computed for every trained model."""

from sklearn.metrics import accuracy_score, classification_report


def evaluate_predictions(y_test, y_pred):
    """Return the weighted F1/precision/recall plus accuracy for one model."""
    report = classification_report(y_test, y_pred, output_dict=True)
    f1 = report['weighted avg']['f1-score']
    precision = report['weighted avg']['precision']
    recall = report['weighted avg']['recall']
    accuracy = accuracy_score(y_test, y_pred)

    return {
        'F1 Score': f1,
        'Precision': precision,
        'Recall': recall,
        'Accuracy': accuracy,
    }


def select_best_model(results_df):
    """Return ``(best_model_name, best_model)`` - the row with the highest F1."""
    best_model_name = results_df.loc[results_df['F1 Score'].idxmax(), 'Model']
    best_model = results_df.loc[
        results_df['Model'] == best_model_name, 'Trained Model'].values[0]
    return best_model_name, best_model