Build a supervised classifier using embedding features.

What you need to know

A pre-trained embedding model has already learned what text means. A linear classifier on frozen embeddings only has to learn where to draw boundaries between your classes in that space — a much easier job, which is why a few hundred examples are often enough.

Logistic regression learns one weight vector per class; the score for a class is a dot product with the text's vector, and softmax turns scores into probabilities. It is fast, hard to overfit with regularisation, and gives calibrated-enough probabilities for a confidence threshold.

Four practices matter more than the choice of model:

• Stratified split — train and test keep the same class proportions, so a rare class is not missing from the test set.
• Balanced class weights — errors on rare classes cost more, so the model does not ignore them.
• Macro metrics — the average of per-class F1, so each class counts equally. Accuracy on a 95/5 split is 95% for a model that never predicts the rare class.
• Deduplicate before splitting — the same message in train and test inflates the score.

1from collections.abc import Callable2import numpy as np3from sklearn.linear_model import LogisticRegression4from sklearn.metrics import accuracy_score, f1_score5from sklearn.model_selection import train_test_split67def build_features(texts: list[str], embed_fn: Callable, batch_size: int = 64) -> np.ndarray:8    vectors = []9    for i in range(0, len(texts), batch_size):10        vectors.extend(embed_fn(texts[i:i + batch_size]))11    X = np.asarray(vectors, dtype=np.float32)12    return X / (np.linalg.norm(X, axis=1, keepdims=True) + 1e-10)1314def train_classifier(texts: list[str], labels: list[str], embed_fn: Callable, seed: int = 42):15    """Train on 80%, report macro F1 and accuracy on the stratified 20% hold-out."""16    texts, labels = map(list, zip(*dict(zip(texts, labels)).items()))   # dedupe texts17    X, y = build_features(texts, embed_fn), np.asarray(labels)18    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, random_state=seed, stratify=y)19    clf = LogisticRegression(max_iter=2000, class_weight="balanced").fit(X_tr, y_tr)20    pred = clf.predict(X_te)21    return clf, {"accuracy": round(accuracy_score(y_te, pred), 3),22                 "macro_f1": round(f1_score(y_te, pred, average="macro"), 3)}2324def predict(clf, texts: list[str], embed_fn: Callable, min_confidence: float = 0.6):25    probs = clf.predict_proba(build_features(texts, embed_fn))26    return [(str(clf.classes_[i]) if p[i] >= min_confidence else "needs_review", round(float(p[i]), 3))27            for p, i in zip(probs, probs.argmax(axis=1))]

The tricky parts:

• dict(zip(texts, labels)) keeps one label per unique text (the last one seen), removing exact duplicates before the split so none leak across it.
• clf.classes_ is the label order behind predict_proba's columns. Indexing it with the argmax gives the label name; assuming the order you passed in is a classic bug.
• max_iter=2000 because the default can stop before convergence on high-dimensional embeddings and print a warning.
• The same build_features is used for training and prediction, so normalisation is identical in both.

Complexity: embedding N texts is N model passes (batched). Training is roughly O(N·d·k) per optimiser iteration — seconds for tens of thousands of rows. Prediction is one embedding plus O(d·k). Memory is O(N·d) for the features.

A real-life example

Synthetic 16-dimensional "embeddings" for three ticket types stand in for a real embedding model, so the run is reproducible:

1rng = np.random.default_rng(0)2topics = np.eye(16)[:3]3names = ["billing", "delivery", "account"]4data = {f"ticket {i}": (names[i % 3], topics[i % 3] + rng.normal(0, 0.25, 16)) for i in range(150)}5fake_embed = lambda ts: [data[t][1] if t in data else EXTRA[t] for t in ts]6EXTRA = {"charged and not delivered": (topics[0] + topics[1]) / 2 + rng.normal(0, 0.05, 16)}78clf, report = train_classifier(list(data), [lab for lab, _ in data.values()], fake_embed)9print(report)                                         # {'accuracy': 0.967, 'macro_f1': 0.967}10print(predict(clf, ["ticket 4", "charged and not delivered"], fake_embed, min_confidence=0.7))11# [('delivery', 0.813), ('needs_review', 0.642)]

Real embeddings are messier than this, and real scores are lower; the mechanics — split, balanced training, macro F1, a confidence gate — are the same.

E-commerce companies route millions of support tickets a month with exactly this: embeddings plus a linear classifier, retrained weekly, with low-confidence tickets sent to a human queue.

Follow-up questions to expect

• "Why not fine-tune the embedding model instead?" — Fine-tuning can add a few points but costs GPU time and a deployment. Try the linear model first; fine-tune only if the gap matters.
• "A class has only 3 examples — what breaks?" — stratify needs at least 2 per class to split, and the model barely learns it. Collect more examples or merge the class.
• "How do you keep it accurate over time?" — Log low-confidence predictions and human corrections, add them to the training set, and retrain on a schedule, watching macro F1 on a fixed hold-out.