Module 6 — Text classification and sentiment analysis

Module 5 gave you vectors. This module trains a classifier on top of them and compares it, honestly, with the TF-IDF baseline from module 3. On the ten thousand English reviews of the corpus, both approaches are strong. Deciding which one goes to production is the point of this module.

The task: 5 stars, not just thumbs up or down

The reviews carry a star rating from 1 to 5. Collapsing them into a binary positive / negative label throws away useful information — a 3-star review is often the most actionable feedback a product team receives. The full task is a five-class classification, which is harder than it looks: the model has to distinguish it's fine, nothing special (3 stars) from it's fine, I like it (4 stars) from language that is very close.

Start by looking at the class distribution.

import collections
counts = collections.Counter(train_labels)
print(counts)
# Counter({5: 5400, 4: 1900, 1: 1300, 3: 800, 2: 600})

Two things stand out. The majority class (5 stars) accounts for more than half of the corpus. Predicting 5 for every review would already reach 54 % accuracy without learning anything. And the minority classes (2 and 3 stars) are ten times less represented — the model will see roughly one 3-star review for every seven 5-star reviews.

Both facts change what you can trust in the results.

Accuracy is not the metric you want

Consider two models. Model A predicts 5 for everything and reaches 54 % accuracy. Model B correctly identifies 90 % of the 1-star reviews and only 60 % of the 5-star ones; its overall accuracy is 63 %. On accuracy alone, Model A looks better than a naive prediction that Model B is actually the useful one.

The confusion matrix shows the whole picture. Read a row as "actual class", a column as "predicted class":

from sklearn.metrics import confusion_matrix, classification_report

preds = model.predict(test_texts)
cm = confusion_matrix(test_labels, preds, labels=[1, 2, 3, 4, 5])
print(cm)
print(classification_report(test_labels, preds, digits=3))

classification_report gives per-class precision and recall. On this task the meaningful numbers are the F1 scores for the two minority classes (1 and 2 stars), because that is where the business impact lives — an angry customer whose complaint is missed churns; a 5-star customer misclassified as 4-star does nothing.

Fine-tuning an encoder

The idea of fine-tuning is to start from a pretrained encoder (module 5), add a small classification head on top of its [CLS] output, and train the whole stack briefly on the labelled reviews. Because the encoder already knows the language, a few thousand examples and two or three epochs are usually enough.

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset

tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=5,
)

def encode(batch):
    return tok(batch["text"], truncation=True, max_length=256, padding="max_length")

train_ds = Dataset.from_dict({"text": train_texts, "label": [l - 1 for l in train_labels]}).map(encode, batched=True)
test_ds  = Dataset.from_dict({"text": test_texts,  "label": [l - 1 for l in test_labels ]}).map(encode, batched=True)

args = TrainingArguments(
    output_dir="out", num_train_epochs=3, per_device_train_batch_size=16,
    learning_rate=2e-5, weight_decay=0.01, evaluation_strategy="epoch",
    logging_steps=50,
)

trainer = Trainer(model=model, args=args, train_dataset=train_ds, eval_dataset=test_ds)
trainer.train()

Three choices deserve a word. DistilBERT is picked over BERT because it is 40 % smaller and 60 % faster for less than 3 points of accuracy loss on classification. The learning rate is $2 \times 10^{-5}$, small on purpose: fine-tuning a pretrained model with a rate above $5 \times 10^{-5}$ almost always collapses the accuracy back to a random one, because the delicate weights are overwritten. And the labels are shifted from [1..5] to [0..4] because the classification head assumes 0-indexed classes.

Handling the imbalance

Three techniques, in order of increasing complexity.

Class weights. Pass the inverse frequency of each class into the loss so that a mistake on a 3-star review costs more than a mistake on a 5-star review.

import torch

freqs = torch.tensor([counts[i + 1] for i in range(5)], dtype=torch.float)
weights = (freqs.sum() / (5 * freqs)).to(model.device)
loss_fct = torch.nn.CrossEntropyLoss(weight=weights)

Undersampling. Draw a training set with equal numbers per class. Simple and effective, at the cost of throwing away data — usually fine when the majority class has many thousands of examples.

Focal loss. Downweights easy examples and focuses training on hard ones. Useful when classes are extreme (1 : 100 or more), overkill on 10 : 1 like here.

Whatever the technique, evaluate on the original, imbalanced test set. Sampling changes what the model learns; it must not change what the metric measures.

Head-to-head with the baseline

The interesting comparison is not TF-IDF vs BERT in a vacuum. It is on the exact same split, with the same metric, and with cost accounted for.

Model	Macro F1	Training	Inference	Memory
TF-IDF + LogReg	~0.51	seconds	30 μs / doc	20 MB
DistilBERT fine-tuned	~0.62	15 min on a GPU	10 ms / doc	260 MB
BERT-base fine-tuned	~0.64	40 min on a GPU	25 ms / doc	440 MB

The transformer wins on quality, and the gap is bigger on the minority classes (2 and 3 stars), where the baseline's synonyms problem hurts the most. But the baseline is hundreds of times cheaper, both to train and to serve. On any task where the 10 points of macro F1 do not justify a GPU in production, the baseline stays. Module 10 revisits this trade-off in a decision tree.

Do not compare a fine-tuned model to a not-tuned baseline

It is easy to make BERT look magical by comparing three tuned epochs of it against a LogisticRegression(max_iter=100, C=1) with no hyperparameter search. Tune both, or neither. The honest baseline is TfidfVectorizer(min_df=2, max_df=0.9, ngram_range=(1, 2)) with a small grid over C in logistic regression.

Save the confusion matrix, not just the score

A single F1 number hides which class the model still gets wrong. In production, the confusion matrix is what shows you where to invest — collect more 2-star reviews, rewrite the class definition, add a rule for reviews under 15 words — and where the model has plateaued.

In summary

• Accuracy hides the minority classes; on an imbalanced 5-class problem, use the confusion matrix and macro F1, and evaluate on the original, imbalanced test set.
• Fine-tuning an encoder with a small learning rate (about $2 \times 10^{-5}$) and 2 to 3 epochs is the standard recipe; DistilBERT is a good default when latency matters.
• Handle imbalance with class weights or undersampling; focal loss only pays for extreme ratios.
• The transformer usually wins on quality by several points of macro F1, at hundreds of times the cost; the baseline still ships when the gap is small.

Next module: extracting the entities named in the reviews — product names, places, brands — with a BIO-tagged sequence model.