Start Learning Contact Us
Data Science with Python

Text Classification

Natural Language Processing · Lesson 23 of 30

Text classification assigns a label to a piece of text. Spam filtering, sentiment analysis, routing support tickets and detecting the language of a message are all examples.

The recipe combines the last lesson with the machine learning module: turn text into TF-IDF vectors, then train a classifier.

A small dataset

reviews = [
    ("great battery life and fast delivery", 1),
    ("excellent screen, really happy with it", 1),
    ("works perfectly, great value", 1),
    ("love this laptop, fast and light", 1),
    ("fantastic sound and easy setup", 1),
    ("very happy, excellent quality", 1),
    ("good price and great performance", 1),
    ("amazing camera, love the colours", 1),
    ("terrible battery, died in a week", 0),
    ("screen cracked, very poor quality", 0),
    ("stopped working, waste of money", 0),
    ("slow delivery and broken on arrival", 0),
    ("awful sound, really disappointed", 0),
    ("poor value, would not recommend", 0),
    ("bad performance and terrible support", 0),
    ("disappointed, the camera is awful", 0),
]
texts = [t for t, _ in reviews]
labels = [l for _, l in reviews]
print(len(texts), "reviews,", sum(labels), "positive")

Sixteen reviews are far too few for a real model. Real sentiment datasets have thousands of examples. This one is small so that it runs in a second and you can see every row.

Train and predict

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

texts = ["great battery life and fast delivery", "excellent screen, really happy with it",
         "works perfectly, great value", "love this laptop, fast and light",
         "fantastic sound and easy setup", "very happy, excellent quality",
         "good price and great performance", "amazing camera, love the colours",
         "terrible battery, died in a week", "screen cracked, very poor quality",
         "stopped working, waste of money", "slow delivery and broken on arrival",
         "awful sound, really disappointed", "poor value, would not recommend",
         "bad performance and terrible support", "disappointed, the camera is awful"]
labels = [1] * 8 + [0] * 8

model = make_pipeline(TfidfVectorizer(), LogisticRegression(C=10))
model.fit(texts, labels)

new = ["really happy with the fast delivery", "terrible quality, very disappointed", "the screen is not great"]
for text, p in zip(new, model.predict_proba(new)[:, 1]):
    print(f"{p:.2f} positive  |  {text}")

The first two sentences are classified clearly. The third is close to 0.5: the model has seen “great” only in positive reviews, and a bag of words cannot tell that “not” reverses it, so it does not recognise the sentence as negative.

C controls regularisation. A larger value lets the model fit the training data more closely, which makes its probabilities more decisive on this tiny dataset.

Which words matter

A linear model has one coefficient per word. Large positive coefficients push towards the positive class.

import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

texts = ["great battery life and fast delivery", "excellent screen, really happy with it",
         "works perfectly, great value", "love this laptop, fast and light",
         "fantastic sound and easy setup", "very happy, excellent quality",
         "good price and great performance", "amazing camera, love the colours",
         "terrible battery, died in a week", "screen cracked, very poor quality",
         "stopped working, waste of money", "slow delivery and broken on arrival",
         "awful sound, really disappointed", "poor value, would not recommend",
         "bad performance and terrible support", "disappointed, the camera is awful"]
labels = [1] * 8 + [0] * 8

vectorizer = TfidfVectorizer()
X = vectorizer.fit_transform(texts)
classifier = LogisticRegression().fit(X, labels)

weights = pd.Series(classifier.coef_[0], index=vectorizer.get_feature_names_out()).sort_values()
print("most negative:", list(weights.head(5).index))
print("most positive:", list(weights.tail(5).index))

Being able to inspect the model like this is a real advantage of simple methods.

Naive Bayes

Naive Bayes is a classic text classifier. It uses the frequency of each word in each class to calculate the probability of a class given the words, and assumes that words occur independently of one another. The assumption is false, yet the method works well and trains almost instantly.

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline

texts = ["great battery life and fast delivery", "excellent screen, really happy with it",
         "works perfectly, great value", "love this laptop, fast and light",
         "terrible battery, died in a week", "screen cracked, very poor quality",
         "stopped working, waste of money", "awful sound, really disappointed"]
labels = [1] * 4 + [0] * 4

model = make_pipeline(CountVectorizer(), MultinomialNB()).fit(texts, labels)
print(model.predict(["great value and excellent battery", "poor quality, stopped working"]))

Evaluating a text classifier

Everything from the model evaluation lesson applies:

  • Hold out a test set, or use cross-validation.
  • Look at precision and recall for each class, not only accuracy.
  • Check class balance. If 95% of messages are not spam, accuracy says little.
  • Read the examples the model gets wrong. They show what it has not learned, such as negation or sarcasm.

When to use what

Approach Use it when
TF-IDF with a linear model You have labelled examples and need something fast and explainable. A strong baseline
A fine-tuned transformer You need the best accuracy and have the data and computing power
A large language model with a prompt You have few or no labelled examples. Covered in the prompt engineering course

Practice

Add ngram_range=(1, 2) to the TfidfVectorizer and add two training reviews that contain “not great” and “not good” with the negative label. Does the third test sentence change?