From abab284e2e785092b045bb7d6ee365650fc3476b Mon Sep 17 00:00:00 2001
From: Maciej Tyczynski <mactyc1@st.amu.edu.pl>
Date: Wed, 28 Jun 2023 13:31:49 +0200
Subject: [PATCH] added parameters and mlflow to dockerfile

---
 Dockerfile |  1 +
 zad1.py    | 13 ++++++++++---
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 5fa0bc5..e87bc9f 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -8,5 +8,6 @@ RUN pip install scikit-learn
 RUN pip install datasets
 RUN pip install torch
 RUN pip install seaborn
+RUN pip install mlflow
 COPY ./zad1.py ./
 CMD ["python3", 'zad1.py']
\ No newline at end of file
diff --git a/zad1.py b/zad1.py
index b3d2a25..59c8906 100644
--- a/zad1.py
+++ b/zad1.py
@@ -6,6 +6,13 @@ import numpy as np
 import logging
 
 
+import argparse
+
+parser = argparse.ArgumentParser(description='IUM script')
+parser.add_argument('--num_epochs', type=int, default=10, help='Number of epochs')
+parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+parser.add_argument('--alpha', type=float, default=0.001, help='Learning rate')
+args = parser.parse_args()
 logging.basicConfig(level=logging.WARN)
 logger = logging.getLogger(__name__)
 
@@ -98,9 +105,9 @@ model = TabularModel(input_dim, hidden_dim, output_dim)
 criterion = nn.CrossEntropyLoss()
 optimizer = torch.optim.Adam(model.parameters())
 
-num_epochs = 10
-lr = 0.01
-alpha = 0.01
+num_epochs = args.num_epochs
+lr = args.lr
+alpha = args.alpha
 model = TabularModel(input_dim=len(wine_train.columns)-1, hidden_dim=hidden_dim, output_dim=output_dim)
 criterion = nn.CrossEntropyLoss()
 optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=alpha)