1
$\begingroup$

I am working on loan default data and my model is not able to make accurate predictions on the test set because the the default percentage on the test set is very different from that of the training set. Is there any way to ensure that the Sklearn train_test_split function splits the data in such a way that the overall behaviour is not changed . Thanks in advance

Here is the data in a tabular format

Here is the data in a tabular format

$\endgroup$
2
  • $\begingroup$ Why are the percentages not the same, just some bad luck in the way the software did the split? $\endgroup$
    – Dave
    Nov 30, 2022 at 7:25
  • $\begingroup$ It can't be bad luck because then it would be random and would not happen for all the values as we see above. At least for some attributes the training and test percentages should have been the same. $\endgroup$
    – J.Sriram
    Nov 30, 2022 at 11:43

1 Answer 1

0
$\begingroup$

Assuming you want to keep the distributions of the different categories of a certain variable in both test and train, try using the stratify parameter.

I'll suppose that in your case, you want to keep the distributions for the "employee_type" variable with categories like: Accountants, Core staff, drivers, etc. I'd use for this:

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=X['employee_type'])
$\endgroup$
2
  • $\begingroup$ Thanks, I will try this out . Although since this problem is coming for all the categorical variables would it work if I used y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=X) instead of what you suggested . Thanks in advance $\endgroup$
    – J.Sriram
    Nov 30, 2022 at 11:58
  • $\begingroup$ You could try this solution here. Basically to stratify for multiple columns, you'll have to combine all those clases in a new one (concatenating the texts) and stratify by this new col. $\endgroup$
    – ru.mp
    Nov 30, 2022 at 12:37

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.