sklearn model.predict wrong shape after split with kf.split

Question

I try to predict my model of text string using sklearn, with the code below

from sklearn import datasets

news = datasets.load_files("dataset-news", encoding='latin1', categories=categories)

def vectorize_data(data):
    count_vect = CountVectorizer()
    return count_vect.fit_transform(data)

# Gaussian naive Bayes
def gaussian_train(train, target):
    gnb = GaussianNB()
    gnb.fit(train, target)
    return gnb

kf = KFold(n_splits=5)
counter = 1

for train_idx, test_idx in kf.split(news.data):
    print ("%d Fold" % counter)
    train_data = vectorize_data(np.array(news.data)[train_idx])
    test_data = vectorize_data(np.array(news.data)[test_idx])

    print("Gaussian naive Bayes")
    print(train_data.shape)
    print(test_data.shape)
    g_model_train = gaussian_train(train_data.toarray(), news.target[train_idx])
    # predict_data(g_model_fold, test_data.toarray(), target_data)
    # Predict unseen test data based on fitted classifer
    predicted = g_model_fold.predict(test_data.toarray())

From my console

1 Fold
Gaussian naive Bayes
(640, 13477)
(161, 5193)

But then i got

ValueError: operands could not be broadcast together with shapes (161,5193) (14214,) 

How to fix this?

Answer 1

When you convert the text into token counts, the features used should be the same so that the matrix has the same number of columns. One option is to return the countvectorizer from the train data and use it on the test data. So we set up the vectorize_data() functions:

from sklearn import datasets
from sklearn.model_selection import KFold
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import GaussianNB

def vectorize_data(data):
    count_vect = CountVectorizer()
    return count_vect.fit(data)

Use an example dataset:

categories = ['alt.atheism', 'sci.space']

news = datasets.fetch_20newsgroups(categories=categories)

Run the kfold :

kf = KFold(n_splits=5)

for train_idx, test_idx in kf.split(news.data):
    
    cvect = vectorize_data(np.array(news.data)[train_idx])
    train_data = cvect.transform(np.array(news.data)[train_idx])
    test_data = cvect.transform(np.array(news.data)[test_idx])

    print("Gaussian naive Bayes")
    print(train_data.shape)
    print(test_data.shape)
    g_model_train = gaussian_train(train_data.toarray(), news.target[train_idx])
    # predict_data(g_model_fold, test_data.toarray(), target_data)
    # Predict unseen test data based on fitted classifer
    predicted = g_model_train.predict(test_data.toarray())

Output:

Gaussian naive Bayes
(858, 20415)
(215, 20415)
Gaussian naive Bayes
(858, 20019)
(215, 20019)
Gaussian naive Bayes
(858, 20094)
(215, 20094)
Gaussian naive Bayes
(859, 20119)
(214, 20119)
Gaussian naive Bayes
(859, 20207)
(214, 20207)