Stratification of multilabel data

[dokato]

Recently, new Multilabel classification task has been added: #440

Looking at datasets available on HF, they're typically quite large, e.g. >50k examples. What's the best way to train/test split it, or just sample it?

Maybe sth from: http://scikit.ml/stratification.html ?

cc @x-tabdeveloping as you added that in

[x-tabdeveloping]

I have added some stratified subsampling for multilabel data in this PR: #694 You can have a look at the code there.
But in essence I think you can just use scikit-learn's train_test_split with an encoded labels array as the stratify argument.

Something like this:

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MultiLabelBinarizer

X: ArrayLike = [...]
labels: list[list[str]] = [...]
encoded_labels = MultiLabelBinarizer.fit_transform(labels)

X_train, y_train, X_test, y_test = train_test_split(X, labels, stratifiy=encoded_labels)

[x-tabdeveloping]

It's not merged yet, so you haven't overlooked anything :D
I'm not sure how well it works though, a lot of times it throws a warning and doesn't reduce the dataset size, so we might have to look into it.

[dokato]

Yea, because on surface it doesn't seem like such a trivial problem, especially when dealing with narrowly represented classes.

[dokato]

@x-tabdeveloping have a look at quick and dirty attempt of using iterative_train_test_split from this scikit-multilearn library:
#546
It works quite neat and fast, though requires a few hack arounds to make it running with our flow, plus we need to add another dependency for just this one function, LMK what you think...

[x-tabdeveloping]

@dokato Can you make this in a separate PR and then we can discuss it there? I have some thoughts about it

[dokato]

DIscussion moved to: #760