From 04fe4bb719873b573200a3cdae44befa6d8d1bd5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Zar=C4=99ba?= <k.mjayk@gmail.com>
Date: Wed, 10 May 2023 23:54:20 +0200
Subject: [PATCH] ss

---
 MLproject |  7 ++++++-
 train.py  | 14 ++++++++++----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/MLproject b/MLproject
index b13701f..4dea0d0 100644
--- a/MLproject
+++ b/MLproject
@@ -1,7 +1,12 @@
 name: s487187-training
 
+docker_env:
+  image: python:3.11
+  volumes: [ "/root/.cache:/root/.cache" ]
+  user_env_vars: [ "SACRED_IGNORE_GIT" ]
+
 entry_points:
   main:
     parameters:
       epochs: int
-    command: "python train.py --epochs {epochs}"
\ No newline at end of file
+    command: "python train.py --epochs {epochs}"
diff --git a/train.py b/train.py
index 3447d1b..6655a5d 100644
--- a/train.py
+++ b/train.py
@@ -1,8 +1,10 @@
 from sacred import Experiment
-from sacred.observers import MongoObserver, FileStorageObserver
+from sacred.observers import MongoObserver
 import os
 import mlflow
 import mlflow.keras
+from mlflow.models.signature import infer_signature
+from mlflow.models import Model
 
 os.environ["SACRED_NO_GIT"] = "1"
 
@@ -12,12 +14,13 @@ ex.observers.append(MongoObserver(url='mongodb://admin:IUM_2021@172.17.0.1:27017
 @ex.config
 def my_config():
     data_file = 'data.csv'
-    model_file = 'model.h5'
+    model_file = 'model'
     epochs = 10
     batch_size = 32
     test_size = 0.2
     random_state = 42
 
+
 @ex.capture
 def train_model(data_file, model_file, epochs, batch_size, test_size, random_state):
     import pandas as pd
@@ -62,7 +65,10 @@ def train_model(data_file, model_file, epochs, batch_size, test_size, random_sta
         mlflow.log_metric("loss", loss)
         mlflow.log_metric("accuracy", accuracy)
 
-        mlflow.keras.save_model(model, model_file)
+        signature = infer_signature(X_train, model.predict(X_train))
+        input_example = Model.log_input_example(X_train.iloc[0])
+
+        mlflow.keras.log_model(model, model_file, signature=signature, input_example=input_example)
 
         return accuracy  
 
@@ -72,4 +78,4 @@ def run_experiment():
     ex.log_scalar('accuracy', accuracy)  
     ex.add_artifact('model.h5')  
 
-ex.run()
+ex.run()
\ No newline at end of file