Skip to content

Commit b0e80a4

Browse files
Add multinomial Naive Bayes text classification example (#14665)
* Add multinomial Naive Bayes text classification example * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review feedback for Naive Bayes text Classifier * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 423eb91 commit b0e80a4

1 file changed

Lines changed: 202 additions & 0 deletions

File tree

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,202 @@
1+
"""
2+
Naive Bayes text classification using a multinomial event model.
3+
4+
The implementation in this module is intentionally educational and keeps the
5+
logic explicit: token counting, prior probabilities, and posterior scoring in
6+
log-space.
7+
8+
References:
9+
- https://en.wikipedia.org/wiki/Naive_Bayes_classifier
10+
- https://scikit-learn.org/stable/modules/naive_bayes.html
11+
"""
12+
13+
from __future__ import annotations
14+
15+
import re
16+
from collections import Counter, defaultdict
17+
from math import exp, log
18+
19+
20+
class NaiveBayesTextClassifier:
21+
"""
22+
Multinomial Naive Bayes classifier for short text documents.
23+
24+
Args:
25+
alpha: Additive (Laplace) smoothing parameter. Must be greater than 0.
26+
27+
>>> NaiveBayesTextClassifier(alpha=0)
28+
Traceback (most recent call last):
29+
...
30+
ValueError: alpha must be greater than 0.
31+
"""
32+
33+
def __init__(self, alpha: float = 1.0) -> None:
34+
if alpha <= 0:
35+
raise ValueError("alpha must be greater than 0.")
36+
37+
self.alpha = alpha
38+
self.classes_: list[str] = []
39+
self.vocabulary_: set[str] = set()
40+
self.class_document_counts_: Counter[str] = Counter()
41+
self.class_token_counts_: dict[str, Counter[str]] = defaultdict(Counter)
42+
self.class_total_tokens_: Counter[str] = Counter()
43+
self.class_log_prior_: dict[str, float] = {}
44+
self.is_fitted_ = False
45+
46+
@staticmethod
47+
def _tokenize(text: str) -> list[str]:
48+
"""
49+
Split text into lowercase alphanumeric tokens.
50+
51+
>>> NaiveBayesTextClassifier._tokenize("Hello, NLP world!")
52+
['hello', 'nlp', 'world']
53+
"""
54+
return re.findall(r"[a-z0-9']+", text.lower())
55+
56+
def fit(self, texts: list[str], labels: list[str]) -> None:
57+
"""
58+
Fit the classifier from labeled training texts.
59+
60+
>>> model = NaiveBayesTextClassifier()
61+
>>> model.fit(["cheap meds", "project meeting"], ["spam", "ham"])
62+
>>> sorted(model.classes_)
63+
['ham', 'spam']
64+
65+
>>> model.fit(["only one text"], ["ham", "spam"])
66+
Traceback (most recent call last):
67+
...
68+
ValueError: texts and labels must have the same length.
69+
70+
>>> model.fit([], [])
71+
Traceback (most recent call last):
72+
...
73+
ValueError: training data must not be empty.
74+
"""
75+
if not texts:
76+
raise ValueError("training data must not be empty.")
77+
if len(texts) != len(labels):
78+
raise ValueError("texts and labels must have the same length.")
79+
80+
self.classes_ = sorted(set(labels))
81+
self.vocabulary_.clear()
82+
self.class_document_counts_.clear()
83+
self.class_token_counts_ = defaultdict(Counter)
84+
self.class_total_tokens_.clear()
85+
self.class_log_prior_.clear()
86+
87+
for text, label in zip(texts, labels):
88+
if not isinstance(text, str) or not isinstance(label, str):
89+
raise TypeError("texts and labels must contain strings only.")
90+
91+
tokens = self._tokenize(text)
92+
self.class_document_counts_[label] += 1
93+
self.class_token_counts_[label].update(tokens)
94+
self.class_total_tokens_[label] += len(tokens)
95+
self.vocabulary_.update(tokens)
96+
97+
total_documents = len(texts)
98+
self.class_log_prior_ = {
99+
label: log(self.class_document_counts_[label] / total_documents)
100+
for label in self.classes_
101+
}
102+
self.is_fitted_ = True
103+
104+
def predict_proba(self, text: str) -> dict[str, float]:
105+
"""
106+
Return posterior probabilities for every class.
107+
108+
>>> train_texts, train_labels = build_toy_dataset()
109+
>>> model = NaiveBayesTextClassifier()
110+
>>> model.fit(train_texts, train_labels)
111+
>>> probs = model.predict_proba("cheap meds available now")
112+
>>> round(sum(probs.values()), 6)
113+
1.0
114+
>>> probs['spam'] > probs['ham']
115+
True
116+
117+
An empty input text has no tokens, so predictions fall back to class priors.
118+
>>> empty_probs = model.predict_proba("")
119+
>>> round(empty_probs['spam'], 3), round(empty_probs['ham'], 3)
120+
(0.5, 0.5)
121+
122+
>>> NaiveBayesTextClassifier().predict_proba("hello")
123+
Traceback (most recent call last):
124+
...
125+
ValueError: model has not been fitted yet.
126+
"""
127+
if not self.is_fitted_:
128+
raise ValueError("model has not been fitted yet.")
129+
if not isinstance(text, str):
130+
raise TypeError("text must be a string.")
131+
132+
tokens = self._tokenize(text)
133+
vocabulary_size = len(self.vocabulary_)
134+
log_posteriors: dict[str, float] = {}
135+
136+
for label in self.classes_:
137+
log_prob = self.class_log_prior_[label]
138+
token_counts = self.class_token_counts_[label]
139+
denominator = self.class_total_tokens_[label] + self.alpha * vocabulary_size
140+
141+
for token in tokens:
142+
count = token_counts[token]
143+
log_prob += log((count + self.alpha) / denominator)
144+
145+
log_posteriors[label] = log_prob
146+
147+
max_log = max(log_posteriors.values())
148+
exp_scores = {
149+
label: exp(score - max_log) for label, score in log_posteriors.items()
150+
}
151+
normalizer = sum(exp_scores.values())
152+
return {label: score / normalizer for label, score in exp_scores.items()}
153+
154+
def predict(self, text: str) -> str:
155+
"""
156+
Predict the most likely class label for a text.
157+
158+
>>> train_texts, train_labels = build_toy_dataset()
159+
>>> model = NaiveBayesTextClassifier(alpha=1.0)
160+
>>> model.fit(train_texts, train_labels)
161+
>>> model.predict("free cheap meds")
162+
'spam'
163+
>>> model.predict("project meeting schedule")
164+
'ham'
165+
"""
166+
probabilities = self.predict_proba(text)
167+
return max(probabilities, key=lambda label: probabilities[label])
168+
169+
170+
def build_toy_dataset() -> tuple[list[str], list[str]]:
171+
"""
172+
Build a tiny text dataset for examples and quick local testing.
173+
174+
>>> texts, labels = build_toy_dataset()
175+
>>> len(texts), len(labels)
176+
(6, 6)
177+
>>> sorted(set(labels))
178+
['ham', 'spam']
179+
"""
180+
texts = [
181+
"buy cheap meds now",
182+
"cheap meds available online",
183+
"win cash prizes now",
184+
"project meeting schedule attached",
185+
"let us discuss the project timeline",
186+
"team meeting moved to monday",
187+
]
188+
labels = ["spam", "spam", "spam", "ham", "ham", "ham"]
189+
return texts, labels
190+
191+
192+
if __name__ == "__main__":
193+
import doctest
194+
195+
doctest.testmod()
196+
197+
sample_texts, sample_labels = build_toy_dataset()
198+
classifier = NaiveBayesTextClassifier(alpha=1.0)
199+
classifier.fit(sample_texts, sample_labels)
200+
201+
print("Prediction:", classifier.predict("cheap prizes available now"))
202+
print("Prediction:", classifier.predict("team meeting about project timeline"))

0 commit comments

Comments
 (0)