Added function for loading training data in model training script

This function is made so that it can be used to load and have nicely
prepared training data.