refactor: replace custom methods with scikit learn ones

eg train_test_split