added parameters and mlflow to dockerfile

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']

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)