Sklearn text classification model returns single class regardless of actual content

Viewed 45

I am building text classification model. And for some reason it returns me single class regardless of actual text or number of rows. Here's what I am doing:

X = df['text']
y = df['type']
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state = 13)

pipe_mnnb = Pipeline(steps = [('tf', TfidfVectorizer()), ('mnnb', MultinomialNB())])

pgrid_mnnb = {
 'tf__max_features' : [100, 250, 500],
 'tf__stop_words' : ['english', None],
 'tf__ngram_range' : [(1,1),(1,2)],
 'tf__use_idf' : [True, False],
 'mnnb__alpha' : [0.1, 0.5, 1]
}
gs_mnnb = GridSearchCV(pipe_mnnb,pgrid_mnnb,cv=5,n_jobs=-1)
gs_mnnb.fit(X_train, y_train)

Then I use following lines of code to double check results:

preds_mnnb = gs_mnnb.predict(X)
df['preds'] = preds_mnnb

Everything looks ok. Now, I create a new dataframe I want to apply classification model to(I am using dummy data here):

test_data = pd.DataFrame({'text':['abc text',
                           'xyz text',
                           'mnopr text',
                           'ijk text',
                           'rrr text',
                           'xxx text']}
                        )
text
abc text
xyz text
mnopr text
ijk text
rrr text
xxx text

If I use this approach then I get correctly classified results with different classes:

test_data2 = test_data.iloc[:,0]
gs_mnnb.predict(test_data2)

array(['class A', 'class B', 'class C', 'class C', 'class D', 'class B'], dtype='<U30')

But when I dont apply iloc data selection:

gs_mnnb.predict(test_data)

Then I will get only this kind of response, regardless how many rows or what they contain:

array(['class A'], dtype='<U30')

1 Answers

If the model needs to work with DataFrames, one thing I would try is making sure that X and y are DataFrames - currently, they are Series. Change the first two lines of code like this (note the extra square brackets):

X = df[['text']]
y = df[['type']]
Related