Classic ML · Medium · 20 min

Classify email without the cloud

A Naive Bayes classifier over 60 emails of a small business. It trains in milliseconds in your tab, and you can explain every weight.

EMAILS holds 60 short emails (synthetic, in Spanish) labelled cliente (customer) or spam. You will write a multinomial Naive Bayes, the most explainable text classifier there is. Five functions:

tokenize(text): lower case, no accents (normalize('NFD') and drop the diacritics \u0300-\u036f), then the runs of a-z0-9 with at least two characters.

train(docs): returns { classes, vocab, docs }. In classes, per label: { docs, tokens, counts }, where counts is how many times each word appears. vocab is the number of distinct words overall and docs the total number of emails.

logProb(model, label, tokens): log(class docs / docs) plus, for each token, log((count + 1) / (class tokens + vocab)). The + 1 is Laplace smoothing: a new word does not sink the probability to zero.

predict(model, text): { label, pSpam }, with pSpam = 1 / (1 + exp(logCliente - logSpam)) and label equal to 'spam' from 0.5 up.

precisionRecall(preds, labels, positive): { precision, recall } of the positive class, 0 when a denominator is 0.

The preview draws the heaviest words, measures precision and recall leaving one email out at a time, and lets you try an email of your own.

Challenges 0/5

  • tokenize: lower case, no accents, no punctuation, no one-letter words
  • train counts docs, tokens and words per class, plus the whole vocabulary
  • logProb adds the prior and the smoothed logs; a new word never gives -Infinity
  • predict gets a spam email right and, leaving one out, at least 9 in 10
  • precisionRecall for the positive class, and zeros when there are no positives

function tokenize(text) {
  // lower case, no accents, then runs of a-z and 0-9 with at least two characters
  return text.split(' ');
}

function train(docs) {
  // per label: docs, tokens and counts per word; plus the vocabulary size and the total docs
  return { classes: {}, vocab: 0, docs: 0 };
}

function logProb(model, label, tokens) {
  // log(prior) + sum of log((count + 1) / (tokens + vocab)) for each token (Laplace smoothing)
  return 0;
}

function predict(model, text) {
  // pSpam = 1 / (1 + exp(logCliente - logSpam)); the label is 'spam' from 0.5 up
  return { label: 'cliente', pSpam: 0 };
}

function precisionRecall(preds, labels, positive) {
  // precision = tp / (tp + fp), recall = tp / (tp + fn); 0 when a denominator is 0
  return { precision: 0, recall: 0 };
}

const model = train(EMAILS);
console.log(predict(model, 'Has ganado un premio, haz clic aquí para reclamarlo'));
Console output appears here (console.log).

The idea

Naive Bayes counts words. For each class it stores how often every word appears and, when a new email arrives, adds up the logarithms of those frequencies. The class with the highest sum wins. It is called "naive" because it assumes words are independent of each other, and it still works well on short text.

With 60 emails the model trains in under a millisecond and every weight has an explanation: "gana" (win) leans towards spam because it appears in five spam emails and in no customer email. That is the case for starting here rather than with a language model: when it is right you know why, and when it is wrong you know why too.

The trap is measuring it on the same emails you trained on. There it gets almost everything right. The preview measures by leaving one email out at a time, the honest version with so little data. And the 0.5 threshold is not sacred: if you would rather never lose a customer email, raise it and accept a little more spam in the inbox.

Run it on your machine

With Python and scikit-learn, on this exercise's CSV (CC BY 4.0). The same count, with cross-validation:

curl -O https://vorluxai.com/data/labs/correos-pyme.csv
python3 -c "
import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import cross_val_score
df = pd.read_csv('correos-pyme.csv')
nb = make_pipeline(CountVectorizer(strip_accents='unicode'), MultinomialNB())
print(cross_val_score(nb, df.texto, df.etiqueta, cv=10).mean())
"

Check what you learned

1. What does the "+ 1" in `(count + 1) / (tokens + vocab)` do?
2. Why add logarithms instead of multiplying probabilities?
3. The preview measures by "leaving one out". What does that avoid?
4. You raise the spam threshold from 0.5 to 0.8. What happens?

Go deeper: the scikit-learn reference →

This in production, with your data? Let's talk for 15 minutes →