Added training and test sets