From cf74d8cc1c077acc35a36de5c551eef0c094219f Mon Sep 17 00:00:00 2001
From: Wojciech Lidwin <59863785+Halal37@users.noreply.github.com>
Date: Fri, 12 May 2023 15:47:59 +0200
Subject: [PATCH] Fix mlflow

---
 Jenkinsfile_train | 2 +-
 ium_train.py      | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/Jenkinsfile_train b/Jenkinsfile_train
index b943ed8..8b9b078 100644
--- a/Jenkinsfile_train
+++ b/Jenkinsfile_train
@@ -29,7 +29,7 @@ node {
 	testImage.inside{ 
          copyArtifacts filter: 'baltimore_train.csv', projectName: 's487197-create-dataset'
 		 sh "python3 ium_sacred.py -epochs $EPOCHS -lr $LR -validation_split $VALIDATION_SPLIT"
-	     archiveArtifacts artifacts: 'baltimore.zip'
+	     archiveArtifacts artifacts: 'baltimore.zip', 'mlflow_model.pth'
 	 stage('Buil'){
      build job: 's487197-evaluation.eg/master', parameters: [string(name: 'build', value: "${currentBuild.number}")], wait:false
 	 }
diff --git a/ium_train.py b/ium_train.py
index 15605b4..d6d363c 100644
--- a/ium_train.py
+++ b/ium_train.py
@@ -72,12 +72,14 @@ def train_model():
         signature = mlflow.models.signature.infer_signature(train_x, model.predict(x_test))
         if tracking_url_type_store != "file":
 
-            mlflow.sklearn.log_model(model, "wines-model", registered_model_name="ElasticnetWineModel", signature=signature)
+            mlflow.sklearn.log_model(mlf_model, "baltimore_model", registered_model_name="BaltimoreModel", signature=signature)
         else:
-            mlflow.sklearn.log_model(model, "model", signature=signature)
+            mlflow.sklearn.log_model(mlf_model, "model", signature=signature)
         hist = pd.DataFrame(history.history)
         hist['epoch'] = history.epoch
+        mlflow.sklearn.save_model(mlf_model, "mlflow_model.pth")
         model.save('baltimore_model')
+        
         shutil.make_archive('baltimore', 'zip', 'baltimore_model')