Build a Bayesian text classifier in Ruby, then check the arithmetic

A classifier says "spam." Before asking whether it is accurate, ask whether its score is even the score you intended to calculate.

Multinomial Naive Bayes is a good model for that exercise. Its state is a set of counts, and a small example can be checked by hand. This Ruby implementation fits in one downloadable file with its tests. It uses synthetic messages so that the entire calculation is visible; it makes no claim about real inbox performance.

Decide what a word means

The tokenizer lowercases text and extracts runs of ASCII letters. "WIN 123" becomes one token, win. Repeated words remain repeated. There is no stemming, stopword removal, email-header parsing or Unicode normalization.

That deliberately small contract has an obvious limitation: Café becomes caf. It is suitable for explaining the arithmetic on the supplied English fixtures, not for quietly treating multilingual input as correctly normalized text.

Training rejects documents with no tokens. Otherwise an empty training record could increase a class prior without contributing any word evidence. Rejecting it is a policy choice, and the tests make that choice explicit. Prediction, by contrast, accepts empty or entirely unknown text and falls back to class priors.

Count documents separately from tokens

For each label, the model stores a document count, a token count and a count for each word. The vocabulary contains every token seen during training across all labels.

The prior is the fraction of training documents assigned to a label. The likelihood estimate for a word adds alpha to that class's word count, then divides by the class's total token count plus alpha times the vocabulary size. The default alpha is one.

The initializer accepts real numeric alpha values, converts them to Float, then checks that the converted value is finite and positive. That order matters: a huge Integer can become infinity, and a tiny Rational can round to zero. Complex values are rejected before comparison. Tests cover those boundaries and check the hand calculation again with alpha equal to one-half.

This is the multinomial model: three appearances of a word contribute three times, rather than becoming a single present/absent feature. The Stanford information retrieval text explains the model and its independence assumption. The code below makes those choices concrete.

def scores(text)
  raise ArgumentError, "train before predicting" if @documents.empty?
  counts = tokens(text).select { |term| @vocabulary.include?(term) }.tally
  total_documents = @documents.values.sum.to_f
  @documents.keys.sort.to_h do |label|
    score = Math.log(@documents.fetch(label) / total_documents)
    denominator = @totals.fetch(label) + @alpha * @vocabulary.length
    counts.each do |term, frequency|
      probability = (@words.fetch(label)[term] + @alpha) / denominator
      score += frequency * Math.log(probability)
    end
    [label, score]
  end
end

Tokens unseen in the global vocabulary are ignored at prediction time. An alternative would be an explicit unknown-word feature, trained under a corresponding policy. Mixing the two approaches accidentally would change the denominator and the meaning of the counts.

Work through one prediction

Train one ham document containing "meet team" and one spam document containing "win win". There are three vocabulary words. Each class has two tokens and half the documents.

With alpha one, the probability assigned to win is three-fifths in spam and one-fifth in ham. The score for the one-word query is the log of the prior plus the log of that likelihood. The test writes those fractions directly:

def test_hand_calculated_log_scores
  scores = small_model.scores("win")
  assert_in_delta Math.log(0.5) + Math.log(3.0 / 5), scores.fetch("spam"), 1e-12
  assert_in_delta Math.log(0.5) + Math.log(1.0 / 5), scores.fetch("ham"), 1e-12
  assert_equal "spam", small_model.predict("win")
end

The expected values are not generated by calling the classifier again through a different wrapper. They come from the counts in the tiny fixture. That makes this a useful test for a wrong denominator or a missing smoothing term.

A second test repeats win and checks that the difference between the two class scores doubles. This catches an implementation that accidentally converts query tokens into a Set and loses frequency.

Logs preserve comparisons without multiplying tiny numbers

Classification only needs to compare class scores. Adding log probabilities preserves the ordering that multiplying probabilities would produce while avoiding a long chain of very small products.

The implementation groups repeated query words with tally, then multiplies each log likelihood by its frequency. A test scores a document containing twenty thousand copies of win and checks that both scores remain finite. This validates the tested calculation, not unlimited numerical stability for every possible corpus or smoothing parameter.

Predict chooses the largest score. Labels are sorted before scoring so an exact tie has a deterministic outcome: the alphabetically first label wins. That is an arbitrary policy, but it is visible and tested. A real product might instead return an ambiguous result.

The scores are not confidence percentages. Even normalizing them across classes would not establish calibration against real-world outcomes. If the application needs a confidence threshold, evaluate that threshold on representative held-out data rather than displaying a log score with a percent sign.

What the tiny evaluation tells us

The suite trains on four synthetic messages and predicts four different combinations of the same small vocabulary. All four expected labels are recovered.

This is a functional check that the training and prediction paths connect. It is not a meaningful accuracy estimate for spam detection. The vocabulary, class balance and writing style were chosen for the example, and the test messages are very close to the training messages.

For an actual evaluation, hold out messages before tuning tokenization or alpha. Track mistakes by class, including ordinary mail marked as spam. Keep related or duplicate messages from leaking across the split, and test recent data separately when language changes over time.

Those steps determine whether a model is useful. Adding another decimal place to the four-message fixture would not.

Run the complete implementation

Download the classifier and tests. It uses Ruby's Set library and activates Minitest 6.0.6. The recorded runtime is Ruby 3.3.2; install the exact test gem version if required.

ruby example.rb --seed 42

The local run passed 11 tests and 30 assertions. Besides the hand calculation and held-out fixture, it checks alpha conversion boundaries, invalid training, an untrained model, deterministic ties, unknown tokens, repeated words and the long-document score.

Try changing the smoothing denominator or removing query frequency. The corresponding arithmetic tests should fail before you begin arguing about accuracy. For a model this small, understanding a single prediction is an achievable and useful first milestone.