Skip to main content

Module 10 — Project: a document classifier in the course language

Nine modules have built the toolbox. This module builds the product. On the same corpus of ten thousand English customer reviews, we run three classifiers to the same star-rating task, then compare them on quality, latency and cost. The point is not to pick a winner in a vacuum: it is to make the pick with the same evidence a good engineer would show a stakeholder.

The specifics of English you have to respect

Every language brings its own traps. Ours are less severe than French elision or Arabic diacritics, but they still bite.

Contractions. don't, won't, it's, they're. A naive punctuation stripper glues them into dont, wont, its, theyre, and the negation n't disappears. Either keep the apostrophe (standard for word tokenisers) or use a tokenizer that splits contractions properly (don'tdo + n't). Modern subword tokenisers do this on their own.

Capitalisation carries meaning. Apple, Congress, Windows are distinct from apple, congress, windows. Any uncased model discards the signal. For sentiment on reviews the cost is modest; for NER (module 7) it is severe.

Negation scope. English negation is short-range (not good, wasn't great) but sneaky (no problems whatsoever — three negations, positive polarity). A bag-of-words model with 2-grams catches the first case and misses the third. A transformer usually handles both, but it is worth testing on your own reviews.

Building the training and evaluation sets

The single decision that decides the value of every result below is how you split the data.

from sklearn.model_selection import train_test_split

train_val, test = train_test_split(reviews, test_size=0.15, stratify=reviews.label, random_state=42)
train, val = train_test_split(train_val, test_size=0.15, stratify=train_val.label, random_state=42)

print(len(train), len(val), len(test)) # ~7200, ~1300, ~1500

Two rules to enforce. Stratify by class, so that the very rare 3-star reviews are represented in every split. Freeze the test set and never look at it until the very last comparison — every peek costs a decimal in trust.

Approach 1 — TF-IDF plus logistic regression

The baseline from module 3, tuned honestly.

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV

pipe = Pipeline([
("tfidf", TfidfVectorizer(min_df=2, max_df=0.9, ngram_range=(1, 2))),
("clf", LogisticRegression(max_iter=2000, class_weight="balanced")),
])

grid = GridSearchCV(pipe,
{"clf__C": [0.5, 1.0, 2.0, 4.0]},
scoring="f1_macro", cv=3, n_jobs=-1)
grid.fit(train.text, train.label)
print(grid.best_params_, grid.best_score_)

class_weight="balanced" handles the imbalance; the grid searches over regularisation. On the review corpus, expect a macro F1 around 0.50-0.52 on the 5-class task.

Approach 2 — Sentence-BERT embeddings plus a linear head

Precompute one vector per review with a small SBERT (module 5), then fit a linear classifier on the vectors.

from sentence_transformers import SentenceTransformer
from sklearn.linear_model import LogisticRegression

sbert = SentenceTransformer("all-MiniLM-L6-v2")
X_train = sbert.encode(train.text.tolist(), batch_size=64, show_progress_bar=True)
X_val = sbert.encode(val.text.tolist(), batch_size=64)

clf = LogisticRegression(max_iter=2000, C=1.0, class_weight="balanced")
clf.fit(X_train, train.label)
print(clf.score(X_val, val.label))

This is often surprisingly strong: the SBERT vectors carry the meaning TF-IDF misses, and a linear head on top costs nothing to train. Expect macro F1 around 0.55-0.58 on the same split, at a training cost of a few minutes on CPU.

Approach 3 — Fine-tune DistilBERT end-to-end

The recipe from module 6, applied on the same split.

from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
TrainingArguments, Trainer)
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_pandas(train.rename(columns={"label": "labels"})).map(encode, batched=True)
val_ds = Dataset.from_pandas(val.rename(columns={"label": "labels"})).map(encode, batched=True)

args = TrainingArguments(
output_dir="prod-run", num_train_epochs=3, per_device_train_batch_size=16,
learning_rate=2e-5, weight_decay=0.01, warmup_ratio=0.06,
evaluation_strategy="epoch", save_strategy="epoch",
load_best_model_at_end=True, metric_for_best_model="f1",
)

Trainer(model=model, args=args, train_dataset=train_ds, eval_dataset=val_ds,
tokenizer=tok, compute_metrics=compute_macro_f1).train()

Expect macro F1 around 0.60-0.63. The gain over Approach 2 is real but modest; the cost multiplies by an order of magnitude.

The comparison that decides

Once the three models are trained, score them on the held-out test set, once and only once, and measure inference latency and infrastructure cost.

MetricTF-IDF + LogRegSBERT + LogRegDistilBERT fine-tuned
Macro F1 (5 classes)~0.51~0.57~0.62
F1 on 1-star reviews~0.74~0.79~0.83
Training20 s CPU4 min CPU15 min on a T4 GPU
Inference (per doc)30 μs CPU3 ms CPU10 ms CPU / 2 ms GPU
Memory20 MB90 MB (SBERT) + 5 MB260 MB
Model file4 MB90 MB260 MB
Monthly cost @ 1M docs / day$5 CPU$40 CPU$300 CPU or $120 GPU

Numbers are indicative; run them on your own infrastructure. The decision tree flows from the table:

  • If the macro-F1 gap is under 3 points, the baseline ships. Twenty times cheaper, twenty times faster, no GPU dependency.
  • If the gap is 3 to 8 points and the traffic is moderate, SBERT + linear head is the best value: same latency class as the baseline, most of the quality gain.
  • If the gap exceeds 8 points, if the minority classes matter operationally, or if the business case funds a GPU, fine-tune.
  • If none of the above fits and cost dominates, buy a smaller model (distilbert-base over bert-base, quantise to int8, distil further).

Deployment considerations

Latency budget. Under 50 ms per request end-to-end is a common product target. A DistilBERT on CPU with batching of 8 fits. On GPU it fits with batching of 64.

Cold start. A DistilBERT model loads in 5 to 8 seconds from disk. For a serverless deployment where every request pays this, prefer a warmed-up container.

Model drift. English on customer reviews evolves — new product names, new slang, new complaint topics. Re-evaluate the model on fresh data every quarter, and re-train yearly. The retraining pipeline should reproduce every step of the training with the same split rules and the same tokenizer version.

Never compare the F1 you saw on training with the F1 someone else saw on their split

Numbers depend on the split, on class balance, on the length filter, on the tokenizer. Publish the exact split with the metric, or the number means nothing to anyone else.

Keep the baseline in the repo

Even after fine-tuning ships, the TF-IDF baseline earns its keep as a smoke test: any time the transformer's F1 drops within 1 point of the baseline, something is wrong in the pipeline (a preprocessing bug, a tokeniser mismatch, a label file corruption). The cheap model catches expensive bugs.

In summary

  • The split decides the numbers: stratify by class, freeze the test set, and never peek at it before the final comparison.
  • Three approaches — TF-IDF + LogReg, SBERT + LogReg, fine-tuned DistilBERT — cover the useful cost/quality trade-off; each earns its place on the right task.
  • The decision rests on the F1 gap on the classes that matter, the traffic volume and the infrastructure available; document this reasoning next to the model in the repo.
  • English specifics — contractions, capitalisation, negation scope — hurt bag-of-words models much more than transformers; keep them in mind when you compare fairly.

Next module: the course recap and the 40-question exam that certifies you can apply this to your own corpus.