a separate predict function would make the code more usuable to test the trained model on future data sets