Move partitioning to a function

I've extract the code doing the partition of the training, validation and testing set to a function.