from network import createNetwork, getOptimizer, getLossFn
from helpers import getData
from tensorflow.keras.utils import to_categorical

# hyperparameters
learningRate = 1e-3
epochs = 8
batchSize = 64

model = createNetwork()
model.compile(optimizer=getOptimizer(learningRate=learningRate),
              loss=getLossFn(),
              metrics=['accuracy'])

digits_inputs, digits_labels = getData()
categorical_labels = to_categorical(digits_labels)

model.fit(digits_inputs, categorical_labels, epochs=epochs, batch_size=batchSize)

# Save the model
model.save("./digit_prediction")