fix

create-dataset.py

@@ -1,6 +1,7 @@
 import pandas as pd
 import os
 from sklearn.model_selection import train_test_split
+import train
 
 CUTOFF = int(os.environ['CUTOFF'])
 adults = pd.read_csv('adult.csv')
@@ -20,3 +21,5 @@ X_test.to_csv('X_test.csv', index=False)
 Y_test.to_csv('Y_test.csv', index=False)
 Y_train.to_csv('Y_train.csv', index=False)
 Y_dev.to_csv('Y_dev.csv', index=False)
+
+train.main()

train.py

@@ -2,23 +2,25 @@ import pandas as pd
 import tensorflow
 from keras.applications.densenet import layers
 
-train_data_x = pd.read_csv('./X_train.csv')
 
-adults_train = train_data_x.copy()
-adults_predict = train_data_x.pop('age')
-normalize = layers.Normalization()
-normalize.adapt(adults_train)
+def main():
+    train_data_x = pd.read_csv('./X_train.csv')
 
-adult_model = tensorflow.keras.Sequential([
+    adults_train = train_data_x.copy()
+    adults_predict = train_data_x.pop('age')
+    normalize = layers.Normalization()
+    normalize.adapt(adults_train)
+
+    adult_model = tensorflow.keras.Sequential([
         normalize,
         layers.Dense(64),
         layers.Dense(1)
-])
+    ])
 
-adult_model.compile(
+    adult_model.compile(
         loss=tensorflow.keras.losses.MeanSquaredError(),
         optimizer=tensorflow.keras.optimizers.Adam())
 
-adult_model.fit(adults_train, adults_predict, epochs=500)
+    adult_model.fit(adults_train, adults_predict, epochs=500)
 
-adult_model.save('model')
+    adult_model.save('model')