from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.naive_bayes import MultinomialNB

categories = [
    'alt.atheism',
    'soc.religion.christian',
    'comp.graphics',
    'sci.med'
]

twenty_train = fetch_20newsgroups(subset='train',
                                  categories=categories,
                                  shuffle=True,
                                  random_state=42)

print(twenty_train.target_names)


print("\n".join(twenty_train.data[0].split('\n')[:3]))

for t in twenty_train.target[:10]:
    print(twenty_train.target_names[t])

count_vect = CountVectorizer()

X_train_counts = count_vect.fit_transform(twenty_train.data)

# convert document frequency into term frequency
# reduces bias towards longer documents
tfidf_transformer = TfidfTransformer()
X_train_tfidf = tfidf_transformer.fit_transform(X_train_counts)

clf = MultinomialNB().fit(X_train_tfidf, twenty_train.target)

docs_new = ['God is love', 'OpenGL on the GPU is fast']
X_new_counts = count_vect.transform(docs_new)
X_new_tfidf = tfidf_transformer.transform(X_new_counts)

predicted = clf.predict(X_new_tfidf)

for doc,category in zip(docs_new, predicted):
    print("{} => {}".format(doc, twenty_train.target_names[category]))


