add training param

train.py

@@ -1,7 +1,7 @@
 import pandas as pd
 from tensorflow import keras
 from tensorflow.keras import layers
-
+import argparse
 
 class RegressionModel:
     def __init__(self, optimizer="adam", loss="mean_squared_error"):
@@ -42,8 +42,12 @@ class RegressionModel:
         self.model.save("model.keras")
 
 
+parser = argparse.ArgumentParser()
+parser.add_argument('--epochs')
+
+args = parser.parse_args()
 model = RegressionModel()
 model.load_data("df_train.csv", "df_test.csv")
-model.train()
+model.train(epochs=args.epochs)
 model.evaluate()
 model.save_model()

training/Jenkinsfile

@@ -6,6 +6,11 @@ pipeline {
               description: 'Which build to use for copying artifacts',
               name: 'BUILD_SELECTOR'
             )
+            string (
+                name: 'EPOCHS',
+                defaultValue: '35',
+                description: 'Epochs to train'
+            )
         }
     stages {
         stage('Checkout repository') {
@@ -28,7 +33,7 @@ pipeline {
 
             steps {
                 sh "chmod +x ./train.py"
-                sh "python ./train.py"
+                sh "python ./train.py --epochs ${params.EPOCHS}"
                 archiveArtifacts artifacts: 'model.keras', onlyIfSuccessful: true
             }
         }