Q:

"Unexpected keyword argument 'axis'" tunning Sckit-Learn's train_test_split function after using Pandas' cut function

I ran into this problem operating over a dataset. My dataset comes in CSV format and has the following structure:
ID,FieldOne,FieldTwo,FieldThree,FieldFour,FieldThree,FieldFour,FieldFive,ToPredict 
1,337,118,4,4.5,4.5,9.65,1,0.92
2,324,107,4,4,4.5,8.87,1,0.76
3,316,104,3,3,3.5,8,1,0.72

The ´ToPredictField´ is a probability that tells me how likely each row is to be picked for some process. That's my class column, and I want to split it into 5 categories: Very_unlikely (<= 0.5), Unlikely (between 0.5 and 0.7), Medium (between 0.7 and 0.8), Likely (between 0.8 and 0.9), Very_likey (> 0.9). I did this by using Pandas cut function like this:
bins = [0, 0.5, 0.7, 0.8, 0.9, 1]
names = ['Very_unlikely', 'Unlikely', 'Medium', 'Likely', 'Very_likely']
dataset['ToPredictField'] = pd.cut(dataset['Chance of Admit '], bins, labels=names)

Now, I try to run train_test_split to split the dataset into 67% train / 33%:
data_X = dataset[['ID','FieldOne','FieldTwo','FieldThree','FieldFour','FieldThree','FieldFour','FieldFive']].values
data_Y = dataset['Chance of Admit '].values

train_X, test_X, train_Y, test_Y = train_test_split(data_X, data_Y, test_size=0.33, random_state=10)

However, I get this error:
/usr/local/lib/python3.6/dist-packages/sklearn/utils/__init__.py in safe_indexing(X, indices)
    214                                    indices.dtype.kind == 'i'):
    215             # This is often substantially faster than X[indices]
--> 216             return X.take(indices, axis=0)
    217         else:
    218             return X[indices]

TypeError: take_nd() got an unexpected keyword argument 'axis'

Do you have any idea what it can be?
Thanks.

A:

I confirm the problem on pandas 0.24.2. To go around this, change
 data_Y = dataset.ToPredictField.cat.codes

That would give you the numeric codes for the categories, which certainly plays nice with sklearn. Or you can simply do
 data_Y = dataset.ToPredictField

but I am not sure how it goes with sklearn.

