split_data function

Split the data into training and test datasets