split_data function

Split data into train and test sets.