from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.neighbors import NearestCentroid

# textes_train, textes_test : listes de descriptions ; y_train, y_test : leur catégorie
for nom, vec in [("comptage", CountVectorizer()), ("tf-idf", TfidfVectorizer())]:
    X_train = vec.fit_transform(textes_train)       # les poids sont appris sur l'apprentissage seul
    X_test = vec.transform(textes_test)
    modele = NearestCentroid().fit(X_train, y_train)
    print(nom, modele.score(X_test, y_test))

# poids les plus lourds d'une catégorie
mots = vec.get_feature_names_out()
print(mots[X_train[y_train == "Compteurs"].mean(axis=0).A1.argsort()[-6:]])