diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 29a9ffadf36755cd0eb45146c42c79c5aba6d61b..265fa2d60952eda03de0a374332103ab41a8c238 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -19,6 +19,7 @@ train_model: - pip install numpy pandas scikit-learn mlflow # Train the model and log to MLflow + - mlflow server --backend-store-uri sqlite:///mlflow.db --default-artifact-root ./mlruns --host 127.0.0.1 --port 5000 & - python model/train_model.py only: diff --git a/model/train_model.py b/model/train_model.py index 69f86cb612316338d60fd40449a40fc598f03199..9b684153904a2114c3d3cfcfa8d58164ee547487 100644 --- a/model/train_model.py +++ b/model/train_model.py @@ -7,6 +7,7 @@ from sklearn.linear_model import LogisticRegression from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score import mlflow import mlflow.sklearn +from mlflow.tracking import MlflowClient import os @@ -35,6 +36,9 @@ ensemble_model = VotingClassifier( voting='soft' ) +mlflow.set_tracking_uri("http://127.0.0.1:5000") +mlflow.set_experiment("Customer Churn Model") + # Start MLflow run with mlflow.start_run(): # Train the model