added splitting into train/test sets