diff --git a/.gitignore b/.gitignore index 62efb6d..7fe5334 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,9 @@ /data /venv */__pycache__ +/models + +models/ +*.pkl +*.joblib +mlflow.db \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 974177c..a9cfd15 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,9 @@ numpy scipy matplotlib seaborn -pytest \ No newline at end of file +pytest +joblib +mlflow +fastapi +uvicorn +flake8 \ No newline at end of file diff --git a/src/train.py b/src/train.py index 72663ba..318d07f 100644 --- a/src/train.py +++ b/src/train.py @@ -1,6 +1,8 @@ import pandas as pd import joblib from pathlib import Path +import mlflow +import mlflow.sklearn from sklearn.model_selection import train_test_split from sklearn.pipeline import Pipeline @@ -62,33 +64,60 @@ def build_pipeline(X: pd.DataFrame) -> Pipeline: ) def main(): - print("Loading data...") - df = load_data(DATA_PATH) + mlflow.set_experiment("Credit Risk Model") - print("Preparing features and target...") - X = df.drop(columns=[TARGET_COL] + DROP_COLS, errors="ignore") - y = df[TARGET_COL] + with mlflow.start_run(): + print("Loading data...") + df = load_data(DATA_PATH) - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42, stratify=y - ) + print("Preparing features and target...") + X = df.drop(columns=[TARGET_COL] + DROP_COLS, errors="ignore") + y = df[TARGET_COL] + + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42, stratify=y + ) + + print("Building model pipeline...") + pipeline = build_pipeline(X_train) + + # Log parameters + mlflow.log_param("model_type", "LogisticRegression") + mlflow.log_param("max_iter", 1000) + mlflow.log_param("class_weight", "balanced") + mlflow.log_param("random_state", 42) + mlflow.log_param("test_size", 0.2) + + print("Training model...") + pipeline.fit(X_train, y_train) + + print("Evaluating model...") + y_pred = pipeline.predict(X_test) + y_proba = pipeline.predict_proba(X_test)[:, 1] + + report = classification_report(y_test, y_pred, output_dict=True) + roc_auc = roc_auc_score(y_test, y_proba) - print("Building model pipeline...") - pipeline = build_pipeline(X_train) + print(classification_report(y_test, y_pred)) + print("ROC-AUC:", roc_auc) - print("Training model...") - pipeline.fit(X_train, y_train) + # Log metrics + mlflow.log_metric("accuracy", report["accuracy"]) + mlflow.log_metric("precision", report["weighted avg"]["precision"]) + mlflow.log_metric("recall", report["weighted avg"]["recall"]) + mlflow.log_metric("f1_score", report["weighted avg"]["f1-score"]) + mlflow.log_metric("roc_auc", roc_auc) - print("Evaluating model...") - y_pred = pipeline.predict(X_test) - y_proba = pipeline.predict_proba(X_test)[:, 1] + # Log model + mlflow.sklearn.log_model(pipeline, "model") - print(classification_report(y_test, y_pred)) - print("ROC-AUC:", roc_auc_score(y_test, y_proba)) + model_file = MODEL_PATH / "credit_risk_logistic.pkl" + joblib.dump(pipeline, model_file) + print(f"Model saved to {model_file}") - model_file = MODEL_PATH / "credit_risk_logistic.pkl" - joblib.dump(pipeline, model_file) - print(f"Model saved to {model_file}") + # Register the model + model_uri = f"runs:/{mlflow.active_run().info.run_id}/model" + mlflow.register_model(model_uri, "CreditRiskLogisticRegression") if __name__ == "__main__": main()