# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
from sklearn.datasets import fetch_20newsgroups


classes = ['sci.med', 'sci.space', 'alt.atheism']
D = fetch_20newsgroups(categories=classes)

print('Text:')
print(20 * '-')
print(D.data[0])
print()
print('Has the class:')
print(20 * '-')
print(classes[D.target[0] - 1])

# +
from sklearn.model_selection import train_test_split


X_train, X_test, y_train, y_test = train_test_split(D.data, D.target, test_size=.1, random_state=1)
# -

# ### Preprocessing inputs

# +
from sklearn.feature_extraction.text import TfidfVectorizer


vectorizer = TfidfVectorizer()
X_train = vectorizer.fit_transform(X_train).toarray()
X_test = vectorizer.transform(X_test).toarray()
# -

# ### Testing model

# +
from sklearn.naive_bayes import MultinomialNB


model = MultinomialNB(alpha=0.1)
model.fit(X_train, y_train)

# +
from sklearn.model_selection import cross_val_score


print('Score on train set:', cross_val_score(model, X_train, y_train, cv=3).mean())
print('Score on test set:', cross_val_score(model, X_test, y_test, cv=3).mean())


# -

# ### Testing by hand

def predict(text):
    return classes[model.predict(vectorizer.transform([
        text,
    ]).toarray())[0] - 1]


# +
dd = [
    'Are antioxidants healthy?',
    'I want to find double stars.',
    'It\'s proving that god doesn\'t exists!'
]

for d in dd:
    print(d, '--->', predict(d))

StarcoderdataJupyterScriptsDedupFiltered