Skip to main content

Naive Bayes with scikit-learn

MultinomialNB models non-negative event counts and is a classic baseline for bag-of-words text. The vectorizer and classifier belong in one pipeline so the vocabulary is learned from training documents only.

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline

documents = [
"team confirms project meeting tomorrow",
"please review the attached project notes",
"lunch meeting moved to friday",
"the release checklist is ready",
"can we discuss the report today",
"project review starts at ten",
"agenda for tomorrow's team meeting",
"thanks for updating the report",
"win cash now claim your prize",
"exclusive offer click for free reward",
"urgent winner claim cash today",
"free voucher limited time offer",
"claim your prize by clicking now",
"winner selected receive a reward",
"cash bonus offer expires today",
"free prize waiting for you",
]
labels = ["ham"] * 8 + ["spam"] * 8

X_train, X_test, y_train, y_test = train_test_split(
documents, labels, test_size=0.25, stratify=labels, random_state=42
)
model = make_pipeline(
CountVectorizer(),
MultinomialNB(alpha=1.0),
)
model.fit(X_train, y_train)
y_pred = model.predict(X_test)

print(confusion_matrix(y_test, y_pred, labels=["ham", "spam"]))
print(classification_report(y_test, y_pred, digits=3))
print(model.predict(["free cash prize"])[0])

alpha=1.0 applies Laplace smoothing, preventing an unseen training count from forcing an entire class score to zero. The example is deliberately small and only demonstrates the API; a real spam system needs representative, time-aware data, adversarial monitoring, and metrics tied to the harm from missed spam and blocked legitimate mail.

Choose the variant that matches the feature distribution: GaussianNB for class-conditional continuous features modeled as Gaussian, MultinomialNB for non-negative counts or count-like values, BernoulliNB for binary indicators, and CategoricalNB for discrete categories encoded per feature. Scaling is not a universal Naive Bayes requirement and can violate the non-negative input contract of MultinomialNB.

Naive Bayes can classify well while producing poorly calibrated probabilities. Validate calibration before treating predict_proba as risk rather than as a ranking score.