Bag-of-Words
Naive-Bayes
NLTK
Text Classification
Machine Learning

Implementing Bag-of-Words Naive-Bayes classifier in NLTK

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

A bag-of-words Naive Bayes classifier is one of the simplest useful text classification pipelines. The idea is straightforward: convert each document into a set of word-based features, then let Naive Bayes learn which words are associated with each label. NLTK makes this pattern easy to implement, which is why it remains a good teaching and prototyping baseline.

Build Features from Words

The first step is turning text into features. In a simple bag-of-words model, the features usually answer questions like:

  • does this word appear
  • how often does this word appear

NLTK's built-in Naive Bayes classifier works naturally with dictionary-style feature mappings, so a common pattern is to use boolean word-presence features.

A Small Working Example

python
1import nltk
2from nltk import word_tokenize
3
4nltk.download("punkt")
5
6documents = [
7    ("I love this movie", "pos"),
8    ("This film was fantastic", "pos"),
9    ("I hated this movie", "neg"),
10    ("This film was awful", "neg"),
11]
12
13all_words = []
14for text, _ in documents:
15    all_words.extend(word_tokenize(text.lower()))
16
17word_features = sorted(set(all_words))
18
19
20def document_features(text):
21    words = set(word_tokenize(text.lower()))
22    return {f"contains({word})": (word in words) for word in word_features}
23
24
25feature_sets = [(document_features(text), label) for text, label in documents]
26train_set = feature_sets
27
28classifier = nltk.NaiveBayesClassifier.train(train_set)
29
30print(classifier.classify(document_features("fantastic movie")))
31print(classifier.classify(document_features("awful film")))
32classifier.show_most_informative_features(5)

This is already a full bag-of-words Naive Bayes classifier. The vocabulary becomes the feature space, and each document is represented by presence or absence of those words.

Why Naive Bayes Works Reasonably Well

Naive Bayes assumes the features are conditionally independent given the class. That assumption is obviously false for natural language, because words influence each other heavily. Even so, the model often performs well on small and medium text problems because many classification tasks are dominated by informative tokens.

It is especially strong as a baseline because:

  • it trains quickly
  • it is easy to inspect
  • it works with small datasets
  • it gives interpretable feature importance

That makes it ideal for learning pipelines and early prototypes.

Improve the Baseline Carefully

A few practical improvements usually help:

  • lowercase consistently
  • remove obvious punctuation tokens
  • restrict the vocabulary to informative words
  • split data into train and test sets

You can also switch from boolean presence features to count-based features, but the simple presence version is often the easiest place to start when learning the classifier API.

For example, keeping only the most frequent words can reduce noise in larger datasets:

python
1from nltk import FreqDist
2
3freq_dist = FreqDist(all_words)
4word_features = [word for word, _ in freq_dist.most_common(1000)]

Then rebuild the document feature dictionaries from that trimmed vocabulary.

You should also evaluate on held-out data:

python
1train_set = feature_sets[:3]
2test_set = feature_sets[3:]
3
4classifier = nltk.NaiveBayesClassifier.train(train_set)
5print(nltk.classify.accuracy(classifier, test_set))

The tiny sample above is too small for meaningful performance, but the pattern is the right one.

Common Pitfalls

  • Treating tokenization as unimportant when it directly defines the feature space.
  • Building the vocabulary from both training and test data, which leaks information.
  • Using every token blindly and letting noise dominate the classifier.
  • Expecting Naive Bayes to capture complex context or word order on its own.
  • Judging the model only on training predictions instead of held-out evaluation.

Summary

  • In NLTK, bag-of-words Naive Bayes is built from dictionary-style word features.
  • Boolean word-presence features are a simple and effective starting point.
  • NLTK's NaiveBayesClassifier makes the training step straightforward.
  • Careful tokenization, vocabulary selection, and train-test separation matter.
  • This model is best treated as a fast, interpretable baseline rather than as the final word in text classification.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.