Multinomial Naive Bayes

This module gives a pure multinomial Naive Bayes classifier over String features and labels, using Lean's HashMap for the fitted count tables. It is not a tensor-indexed neural model like most of NN/Spec/Models/*; it is a non-neural baseline that keeps the training and prediction semantics explicit.

Probabilities are computed in log space (via MathFunctions.log) to avoid underflow.

Ecosystem note: PyTorch does not provide a Naive Bayes classifier in torch.nn; the closest ecosystem analogue is scikit-learn’s MultinomialNB.

What "training" means here

Naive Bayes is a counting model: training is just collecting label and feature counts from the dataset. The API keeps fitting and inference separate, so examples can show exactly where counts are learned and where predictions are made.

The API exposes an explicit fit step that produces a Model, plus:

• predict for inference using the fitted counts
• negLogLikelihood as a standard training objective (useful for evaluation/comparison)

Implementation status

No neural API builder implements this model, and no model-correctness theorem is proved here.

inductive NaiveBayes.Error : Type

Inputs for which the fitted count tables do not define the requested posterior.

• emptyTrainingData : Error

  Fitting needs at least one labeled example. Feature bags themselves may be empty.

• noLabels : Error

  A manually constructed model has no class labels.

• unknownLabel (label : String) : Error

  Evaluation requested a class outside the fitted label support.

• unknownFeature (feature : String) : Error

  Inference supplied a token outside the fitted vocabulary.

@[instance_reducible]

instance NaiveBayes.instReprError : Repr Error

def NaiveBayes.instReprError.repr : Error → ℕ → Std.Format

@[instance_reducible]

instance NaiveBayes.instBEqError : BEq Error

def NaiveBayes.instBEqError.beq : Error → Error → Bool

structure NaiveBayes.Example : Type

One training example: a bag-of-words feature multiset and a class label.

• features : Array String

  The feature tokens of the example; duplicates are kept so it acts as a multiset.

• label : String

  The class label of the example.

def NaiveBayes.instReprExample.repr : Example → ℕ → Std.Format

@[instance_reducible]

instance NaiveBayes.instReprExample : Repr Example

Fitted model

Model stores the counts and some precomputed bookkeeping derived from the dataset. Nothing here depends on the scalar type α; we only need α when we turn counts into smoothed probabilities (log-space scores).

structure NaiveBayes.Model : Type

Fitted multinomial Naive Bayes model.

This stores raw counts plus a little derived bookkeeping (labels, vocab, totalExamples). Scoring functions turn these counts into Laplace-smoothed log probabilities on demand.

• labelCounts : Std.HashMap String ℕ

  Number of training examples seen for each label.

• featureCounts : Std.HashMap String (Std.HashMap String ℕ)

  For each label, the number of occurrences of each feature token.

• totalCounts : Std.HashMap String ℕ

  For each label, the total number of feature tokens across its examples.

• labels : Array String

  The distinct labels seen during fitting.

• vocab : Array String

  The distinct feature tokens seen during fitting.

• totalExamples : ℕ

  The total number of training examples.

def NaiveBayes.fit (data : Array Example) : Except Error Model

Collect counts from a nonempty training set.

An empty bag of words is legal: its prediction depends only on the class priors. An empty training set is different, since it supplies no classes and therefore defines no posterior.

Scoring and prediction

We use the standard multinomial NB scoring rule (with Laplace smoothing):

• prior: (count(label)+1) / (N + nLabels)
• conditional: (count(feature,label)+1) / (totalFeatures(label) + vocabularySize)

Scores are in log space. For prediction we only need relative ordering.

def NaiveBayes.score {α : Type} [TorchLean.Storage α] [Context α] (m : Model) (input : Array String) (label : String) : Except Error α

Compute a log score for a fitted class and a bag of known feature tokens.

def NaiveBayes.predict (m : Model) (input : Array String) (α : Type) [TorchLean.Storage α] [Context α] : Except Error String

Predict a fitted label, reporting an unfitted model or an unknown token explicitly.

The empty string remains a valid class label; it is never used to signal failure.

Training objective (negative log-likelihood)

This is the standard objective used to evaluate NB models:

- Σ log P(y_i | x_i)

Even though we don't optimize it with gradients (NB training is closed-form counting), having this objective is useful for:

• checking improvements (smoothing choices, feature engineering)
• comparing NB against other baselines
• unit tests / runtime checks

def NaiveBayes.negLogLikelihood {α : Type} [TorchLean.Storage α] [Context α] (m : Model) (data : Array Example) : Except Error α

Negative log-likelihood over fitted classes and known feature tokens.

Every target label belongs to the same class support used in the normalization. Unknown targets are rejected, rather than producing a value that can be negative. Empty evaluation data has loss zero for a model with at least one fitted class.