bpmatt/wildroot-mail-v1

Wildroot Mail v1

Two small models that sort email on the user's own machine. Both are BAAI/bge-small-en-v1.5 (33.4M parameters) fine-tuned end to end and exported to float32 ONNX. They run on CPU with ONNX Runtime and need no model server.

File	Answers	Outputs
student.onnx	Lane, reply urgency, personal	category [n, 5]: probabilities for needs_reply, updates, promotional, sales, spam. urgency [n, 5]: probabilities for 1 to 5. personal [n]: probability. Its action output is unused.
action.onnx	Needs action today	action [n]: score from 0 to 1

The other files:

File	What
vocab.txt, tokenizer.json	The base model's uncased WordPiece tokenizer
config.json	File hashes, token limit and rule thresholds
engine.mjs	Tokenizer and rules in JavaScript, no dependencies
replay.mjs, recorded.jsonl	929 public emails with the expected tokens, answers and decisions
RECEIPT.json	Source commit, hashes, runtime versions and test results

Input

Both models take input_ids, attention_mask and token_type_ids (int64, [n, tokens], at most 256 tokens) and facts (float32, [n, 2]): whether the user has written to the sender before, and whether the headers declare bulk mail.

The text must be in the form the models were trained on:

From: Name <address>
To: up to three addresses (+N more)
Cc: up to three addresses (+N more)
Subject: subject, or (no subject)
Received: 2001-11-14T17:28:02.000Z
Has unsubscribe header: yes or no
Mailing list: list id
Declared automated or bulk: yes
User has previously replied in this thread: yes or no

Body:
plain text body

Omit the To, Cc, Mailing list and Declared automated or bulk lines when they do not apply. Every row of recorded.jsonl has a state field in this form.

Usage

JavaScript, with onnxruntime-node 1.30.0:

import fs from 'node:fs';
import * as ort from 'onnxruntime-node';
import { decide, isBulk, factInputs, loadVocab, tokenize, urgentVerdict } from './engine.mjs';

const config = JSON.parse(fs.readFileSync('config.json', 'utf8'));
const vocab = loadVocab(fs.readFileSync('vocab.txt', 'utf8'));
const student = await ort.InferenceSession.create('student.onnx');
const action = await ort.InferenceSession.create('action.onnx');

const facts = { bulk: isBulk(state), acceptsPosts: false, senderKnown: true };
const ids = tokenize(state, vocab, config.maxTokens);
const shape = [1, ids.length];
const feeds = {
  input_ids: new ort.Tensor('int64', BigInt64Array.from(ids, BigInt), shape),
  attention_mask: new ort.Tensor('int64', new BigInt64Array(ids.length).fill(1n), shape),
  token_type_ids: new ort.Tensor('int64', new BigInt64Array(ids.length), shape),
  facts: new ort.Tensor('float32', Float32Array.from(factInputs(facts)), [1, 2]),
};

const out = await student.run(feeds, ['category', 'urgency', 'personal']);
const answers = { category: [...out.category.data], urgency: [...out.urgency.data], personal: out.personal.data[0] };
const decision = decide(answers, facts, config.rules);
const score = (await action.run(feeds)).action.data[0];
const verdict = urgentVerdict(score, facts, config.rules);

Python, with onnxruntime 1.30.0 and tokenizers 0.23.2:

import numpy as np, onnxruntime as ort
from tokenizers import Tokenizer

tok = Tokenizer.from_file('tokenizer.json')
tok.enable_truncation(max_length=256)
tok.no_padding()
session = ort.InferenceSession('student.onnx', providers=['CPUExecutionProvider'])

e = tok.encode(state)
feeds = {'input_ids': np.array([e.ids], dtype=np.int64),
         'attention_mask': np.array([e.attention_mask], dtype=np.int64),
         'token_type_ids': np.array([e.type_ids], dtype=np.int64),
         'facts': np.array([[1.0, 0.0]], dtype=np.float32)}  # sender known, not bulk
category, urgency, personal = session.run(['category', 'urgency', 'personal'], feeds)

To check a runtime against the recorded answers:

node replay.mjs . onnxruntime-node 2

It exits 1 on any difference.

Rules

decide and urgentVerdict in engine.mjs apply rules the models cannot override. The thresholds are in config.json.

• Bulk mail is never needs_reply, except discussion-list mail that accepts posts with p(needs_reply) of at least 0.8. That exception has not been confirmed on held-out mail.
• promotional from a list that accepts posts becomes updates.
• sales becomes needs_reply when p(needs_reply) is at least 0.2.
• Urgency is capped at 3 for a sender the user has never written to.
• An action score of at least 0.3706 is urgent for a known sender and non-bulk mail, and unverified otherwise.

Results

The test used 546 held-out public emails, scored once. Labels came from two Claude models working blind; a label was kept only where they agreed. The set was chosen to hold more urgent mail than a normal inbox. The reference is Jev 1.14, a hosted model, in its best setup for each measure. Brackets show the 95% interval of the difference.

Measure	This model	Jev 1.14	Difference
Category	90.2%	92.4%	-2.2 [-5.1, 0.6]
Needs-reply precision	87.8%	84.8%	+3.0 [-3.7, 9.6]
Needs-reply recall	90.8%	98.3%	-7.6 [-13.1, -2.5]
Reply urgency, exact	84.4%	83.9%
Personal	96.7%	95.2%
Action, AUC	0.965	0.993	-0.028 [-0.057, -0.004]

The action threshold was chosen on development mail after the test. At that threshold the test set had 23 of 27 urgent emails caught and 43 of 495 others flagged. Those figures were read after the test.

Each model takes about 40 ms per email on 4 CPU threads. Both together use about 460 MiB.

Training

About 20,700 public emails: Enron (80%), phishing samples from phishing_pot, SpamAssassin, and Apache, GNU and Linux kernel mailing lists. A larger language model labelled them. student.onnx was trained for 3 epochs on all tasks at once. action.onnx started from it and was trained on the action task alone, with more urgent examples found by the first model.

Limits

• English only. Text beyond 256 tokens is ignored.
• Most training mail is from Enron, around 2000. Nearly all urgent examples are Enron mail. Modern automated alerts have not been tested.
• All labels are AI judgments.
• It has not been tested on a real personal inbox.
• Lanes are suggestions. Do not hide mail based on them.

Licence

MIT. The base model, BAAI/bge-small-en-v1.5, is MIT.