Ippoboi/mmbert-base-email-classifier

Email Classifier (mmBERT-base ONNX, v9)

A dual-head mmBERT-base classifier for multilingual email category + action prediction, optimized for on-device inference with ONNX Runtime.

Successor to Ippoboi/mmbert-s-email-classifier (v8, mmBERT-small). The two are not interchangeable — v9 changes both the encoder and the input format, and is evaluated on a different, harder test set. See Compared to v8 below.

Model Description

Classifies emails into 6 categories and predicts whether action is required:

Category	Description
PERSONAL	1:1 human communication, social messages, direct correspondence
NEWSLETTER	Subscribed editorial/digest content (curated articles, weekly roundups)
PROMOTIONAL	Marketing pushes, sales, discount offers, product launches
TRANSACTION	Orders, receipts, payments, shipping confirmations
ALERT	Security notices, account warnings, important notifications
SOCIAL	Social network notifications, community updates, reactions

Input format (changed in v9)

v9 feeds the sender to the model and preserves link domains — the two strongest signals v8 was blind to. Boilerplate footers are stripped, and truncation is token-budgeted rather than character-budgeted (v8's 2,000-char budget overflowed the 384-token cap for 37.6% of emails, silently discarding the preserved tail).

From: {display name} <{domain}>
Subject: {subject}

Body: {body, footers stripped, URLs → [URL:domain.com], token-budgeted head+tail}

Feeding a v8-formatted string (no From: line) to this model is out-of-distribution and will degrade accuracy without erroring.

Output format

Single forward pass producing two tensors:

• category_probs: Float32[6] — softmax probabilities per category (argmax = predicted category)
• action_prob: Float32[1] — sigmoid probability of action required

Use the calibrated threshold 0.525 for action_prob, not 0.5. INT8 quantization deflates the action logit; the threshold is swept on the validation set post-quantization and shipped in export_metadata.json as action_threshold_int8. Read it from there rather than hardcoding.

Model Details

Attribute	Value
Base Model	jhu-clsp/mmBERT-base (ModernBERT family, multilingual)
Parameters	~307M
Architecture	mmBERT encoder (22 layers, RoPE + GeGLU + alternating local/global attention) + dual classification heads
Pooling	Mean pooling over last_hidden_state (masked)
ONNX Size	308.6 MB (INT8 dynamic per-channel, encoder layer 11 kept FP32)
Max Sequence	384 tokens
Tokenizer	Gemma 2 BPE (256K vocab) — byte-identical to v8's
Opset	14

Performance

Evaluated on a held-out 902-sample multilingual test set (md5 4743b2b93740e9773b41fa34b22d6884), built to be sender-aware: exact-text deduplicated, capped at ≤5 rows per template, and split into 452 new-sender rows (sender domains that appear nowhere in training) and 450 seen-sender rows.

Metric	INT8 (this artifact)
Category accuracy (overall)	85.92%
— new-sender slice	78.3%
— seen-sender slice	~92%
Category accuracy (English)	84.85%
Category accuracy (French)	86.89%
Action accuracy	91.91%
Action F1 @ 0.525	0.895 (P 0.936 / R 0.857)
Argmax-match vs PyTorch FP32	94.90%

Per-class recall

Class	Recall
ALERT	85.31%
NEWSLETTER	80.41%
PERSONAL	94.97%
PROMOTIONAL	79.67%
SOCIAL	91.57%
TRANSACTION	89.32%

Do not compare 85.92% to v8's 93.15%. They are measured on different test sets. v8's 321-row eval let 73% of test rows share a sender domain with training and did not deduplicate templates, so it rewarded memorization; v9's eval reports generalization separately. The v9 number is lower and more honest — the seen-sender slice (92%) is the closest apples-to-apples comparison, and the 14pp seen/new gap is the real remaining headroom.

The residual errors concentrate in the PROMOTIONAL ↔ NEWSLETTER ↔ ALERT triangle (~50 of 136 test errors), which is reflected in those classes' recall.

Intended Use

• Primary: On-device email triage in multilingual mobile apps (iOS/Android)
• Runtime: ONNX Runtime React Native (default CPU/MLAS execution provider)
• Use case: Prioritizing inbox, filtering noise, surfacing actionable emails — for English- and French-speaking users

How to Use

ONNX Runtime (React Native)

import { InferenceSession, Tensor } from 'onnxruntime-react-native';

const session = await InferenceSession.create('model.onnx');

// Two inputs only — NO token_type_ids (mmBERT has no segment embeddings)
const outputs = await session.run({
  input_ids: inputIdsTensor,           // int64[1, S], S ≤ 384
  attention_mask: attentionMaskTensor, // int64[1, S]
});

const categoryProbs = outputs.category_probs.data; // Float32[6]
const actionProb = outputs.action_prob.data[0];    // Float32

const CATEGORIES = ['ALERT', 'NEWSLETTER', 'PERSONAL', 'PROMOTIONAL', 'SOCIAL', 'TRANSACTION'];
const category = CATEGORIES[categoryProbs.indexOf(Math.max(...categoryProbs))];

// Threshold from export_metadata.json — NOT 0.5
const actionRequired = actionProb > 0.525;

Special tokens (Gemma 2 BPE)

Token	ID
<pad>	0
<eos>	1
<bos>	2
<unk>	3

Sequence wrap: [<bos>, ...content..., <eos>]. There is no [CLS] / [SEP].

Files

File	Size	Description
model.onnx	308.6 MB	INT8 quantized ONNX model
tokenizer.json	32.8 MB	Gemma 2 BPE tokenizer (256K vocab)
tokenizer_config.json	45 KB	Tokenizer configuration
special_tokens_map.json	1 KB	Special token IDs
export_metadata.json	1 KB	Provenance, calibrated action threshold, canonical metrics

Architecture

Input → mmBERT-base Encoder (22 layers, 768 hidden, RoPE + GeGLU)
                       ↓
              Mean-pool over last_hidden_state (masked by attention_mask)
                       ↓
                 ┌─────┴─────┐
                 ↓           ↓
           Category Head  Action Head
           Linear(768→6)  Linear(768→1)
                 ↓           ↓
             softmax      sigmoid

Compared to v8 (mmBERT-small)

	mmBERT-small v8	mmBERT-base v9 (this)
Base	mmBERT-small (~140M)	mmBERT-base (~307M)
Bundle size	135 MB	309 MB
Sender in input	✗	✓ From: line
Link domains	[URL] (domain discarded)	[URL:domain.com]
Footer boilerplate	kept (burns token budget)	stripped
Truncation	1500+500 chars (overflowed the 384-token cap for 37.6% of emails)	token-budgeted head+tail
Eval	321 rows, sender-leaky	902 rows, deduped, sender-disjoint slice
Quantization	INT8 dynamic per-channel	+ encoder layer 11 kept FP32
Action threshold	0.175	0.525
Tokenizer	Gemma 2 BPE	identical

Notes on quantization

INT8 dynamic per-channel quantization via onnxruntime.quantization.quantize_dynamic(weight_type=QInt8, per_channel=True), excluding encoder layer 11, which stays FP32.

A leave-one-out sweep across all 22 encoder layers found layer 11 to be the sole quantization-fragile layer for this checkpoint. Quantizing it costs ~7pp of SOCIAL recall (91.6 → 84.3) and deflates the action logit hard enough to collapse the tuned action threshold to the sweep floor. Excluding it costs +14.3 MB and recovers both, plus ~3.7pp of PyTorch argmax agreement (91.2 → 94.9).

Excluding the classification heads instead (dynamic_pc_fp32_heads) was measured and made no difference — the fragility is in the encoder, not the heads.

The fragile layer is per-seed: do not assume layer 11 for a different checkpoint without re-running the sweep. Static-INT8 calibration (percentile/entropy) remains infeasible on ModernBERT-style graphs in ORT 1.24 due to peak-RAM blowup.

Training data

• Source: Personal Gmail inboxes (anonymized)
• Languages: English, French (joint category × language stratified balance)
• Training rows: 5,008 (after exact-text dedupe and ≤5-per-template capping)
• Class weights: Gentle (max 1.311) — aggressive upweighting has been observed to collapse minority classes under INT8 quantization

Limitations

• Trained on English and French only; may not generalize to other languages despite the multilingual base
• Personal/consumer email patterns; may not generalize to enterprise/corporate email
• The PROMOTIONAL ↔ NEWSLETTER boundary is genuinely fuzzy; it accounts for the largest share of residual errors
• New-sender accuracy (78.3%) trails seen-sender accuracy (~92%) by ~14pp — the model still leans on sender familiarity
• 256K vocab is oversized for an EN+FR-only deployment but is required to use the pretrained mmBERT weights
• Ground-truth labels come from a single LLM labeler, so reported accuracy measures agreement with that labeler, not with end users

License

Apache 2.0