Naive Bayes classifier in Python

naive_bayes_classifier.py

#Import datasets and libraries
from sklearn.datasets import fetch_20newsgroups
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import CountVectorizer

#Category selection to compare different datasets of emails and evaluate how hard is to distinguish those.
the_categories = ['comp.sys.ibm.pc.hardware', 'rec.sport.hockey']
train_emails = fetch_20newsgroups(categories = the_categories, subset = 'train', shuffle = True, random_state = 108)
test_emails = fetch_20newsgroups(categories = the_categories, subset = 'test', shuffle = True, random_state = 108)

#Init and fit counter with a complete list of emails (test + train)
counter = CountVectorizer()
counter.fit(test_emails.data + train_emails.data)

train_counts = counter.transform(train_emails.data)
test_counts = counter.transform(test_emails.data)

#Init Classifier and fit with train data
classifier = MultinomialNB()
classifier.fit(train_counts, train_emails.target)

#Calculate the score on the test dataset
the_score = classifier.score(test_counts, test_emails.target)

print(the_score)