StratifiedKfold over heterogeneous DataFrame

Question

I have a pandas DataFrame which contains string and float columns that needs to be split into balanced slices in order to train a sklearn pipeline.

Ideally I'd use StratifiedKFold over the DataFrame to get smaller chunks of data to cross validate. But it complains that I have unorderable types, like this:

import pandas as pd
from sklearn.cross_validation import StratifiedKFold

dataset = pd.DataFrame(
    [
        {'title': 'Dábale arroz a la zorra el abad', 'size':1.2, 'target': 1},
        {'title': 'Ana lleva al oso la avellana', 'size':1.0, 'target': 1},
        {'title': 'No te enrollé yornetón', 'size':1.4, 'target': 0},
        {'title': 'Acá sólo tito lo saca', 'size':1.4, 'target': 0},
    ])
skfs = StratifiedKFold(dataset, n_folds=2)

>>>  TypeError: unorderable types: str() > float()

There are ways to get folds indices and do slicing over the DataFrame, but I don't think that guarantees that my classes are going to be balanced.

What's the best method to split my DataFrame?

Answer 1

sklearn.cross_validation.StratifiedKFold is deprecated since version 0.18 and will be removed in 0.20. So here is an alternative approach:

from sklearn.model_selection import StratifiedKFold

skf = StratifiedKFold(n_splits=2)
t = dataset.target
for train_index, test_index in skf.split(np.zeros(len(t)), t):
    train = dataset.loc[train_index]
    test = dataset.loc[test_index]

Answer 2

StratifiedKFold requires the number of splits, and the .split() method uses the label's class distribution to stratify the samples. Assuming your label is target, you would:

from sklearn.model_selection import StratifiedKFold
skf = StratifiedKFold(n_splits=2)
X=dataset.drop('target', axis=1)
y=dataset.target
for train_index, test_index in skf.split(X, y):
    X_train, X_test = X.iloc[train_index], X.iloc[test_index]
    y_train, y_test = y.iloc[train_index], y.iloc[test_index]