Text classification with TF-IDF and embeddings

#Why it matters

NLP appeared in 48% of the AI/ML Engineer and 38% of the Data Scientist job descriptions we read. Before reaching for a large language model, strong teams build a cheap baseline. For many classification tasks, a TF-IDF model trained in seconds is as good as anything else and costs nothing to run.

#The concept

Models need numbers, so text must become vectors:

• Bag of words / TF-IDF: one column per word or word pair; values count how often a term appears, down-weighted if it appears in many documents. Sparse, fast and interpretable.

• Embeddings: a neural model maps each text to a dense vector (here 384 numbers) where similar meanings sit close together. They capture paraphrases and synonyms that word counts miss.

Either vector type can feed a classifier such as logistic regression.

#Worked example: SMS spam

The UCI SMS Spam Collection (CC BY 4.0) has 5,572 messages: 4,825 legitimate ("ham") and 747 spam.

import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import f1_score, precision_score, recall_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sentence_transformers import SentenceTransformer

sms = pd.read_csv("SMSSpamCollection", sep="\t", header=None, names=["label", "text"])
X_train, X_test, y_train, y_test = train_test_split(
    sms.text, sms.label, test_size=0.25, random_state=42, stratify=sms.label)

def report(name, pred):
    print(name, "precision %.3f recall %.3f f1 %.3f" % (
        precision_score(y_test, pred, pos_label="spam"),
        recall_score(y_test, pred, pos_label="spam"),
        f1_score(y_test, pred, pos_label="spam")))

tfidf = make_pipeline(
    TfidfVectorizer(ngram_range=(1, 2), min_df=2),
    LogisticRegression(max_iter=1000, class_weight="balanced"),
).fit(X_train, y_train)
report("tf-idf", tfidf.predict(X_test))

encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")   # Apache 2.0, about 90 MB
emb_train = encoder.encode(X_train.tolist(), batch_size=64)
emb_test = encoder.encode(X_test.tolist(), batch_size=64)
embedded = LogisticRegression(max_iter=2000, class_weight="balanced").fit(emb_train, y_train)
report("embeddings", embedded.predict(emb_test))

The TF-IDF model wins, trains in under a second and runs anywhere. Why? Spam is defined by specific surface words and patterns: the largest positive TF-IDF weights are "call", "txt", "free", "text", "reply", "stop", "uk" and "www". A general-purpose embedding model smooths away exactly those details, because it was trained to capture meaning, not phone-number-and-prize patterns.

Embeddings earn their place when meaning matters more than exact words: search, grouping support tickets by topic, matching questions to answers. You will use them for retrieval in the RAG shelf.

#Common mistakes

• Skipping the baseline. Without it you cannot say whether an LLM adds anything for its cost.

• Fitting the vectoriser on all data. Keep it inside the pipeline so the vocabulary comes from training data only.

• Reporting accuracy on imbalanced classes. 87% of messages are ham; report precision, recall and F1 for the class you care about.

• Ignoring errors. Read twenty misclassified messages; they tell you what features are missing.

#Interview angle

"How would you classify support tickets?" Start with a TF-IDF baseline and a clear metric, then try embeddings or an LLM only if the baseline falls short on the cases that matter. Quoting a result like this one, where the simple model won, shows judgement.

#Practice task

1. Print the 20 most negative TF-IDF weights (the strongest ham words).

2. Find five test messages both models got wrong. What do they have in common?

3. Try ngram_range=(1, 1) and character n-grams (analyzer="char_wb", ngram_range=(3, 5)). Which helps?

#Further reading

• UCI SMS Spam Collection (CC BY 4.0)

• scikit-learn: text feature extraction

• Sentence Transformers documentation