Retrieving training and testing datasets

We prepare the dataset for the model and let’s divide df_merged into train and test, as it was in the beginning.