Wildroot Mail v1: two ONNX models, tokenizer, rules and replay

32215fc 3 days ago

5.13 kB

// The Wildroot Mail engine: the tokenizer and rules around the student models. No dependencies.
// The host runs the ONNX files. This file turns mail into model inputs and model answers into a decision.

export const CATEGORIES = ['needs_reply', 'updates', 'promotional', 'sales', 'spam'];
const BODY_MARKER = /(?:^|\n)Body:\s*\n/;
const BULK_HEADER = /(?:^|\n)(?:Has unsubscribe header: yes|Mailing list: [^\n]+|Declared automated or bulk: yes)(?=\n|$)/;
const CJK = [[0x4E00, 0x9FFF], [0x3400, 0x4DBF], [0x20000, 0x2A6DF], [0x2A700, 0x2B73F], [0x2B740, 0x2B81F],
  [0x2B920, 0x2CEAF], [0xF900, 0xFAFF], [0x2F800, 0x2FA1F]];
const WORD_LIMIT = 100;

const isCjk = (code) => CJK.some(([low, high]) => code >= low && code <= high);
const isSpace = (c) => c === '\t' || c === '\n' || c === '\r' || /\p{White_Space}/u.test(c);
const isControl = (c) => c !== '\t' && c !== '\n' && c !== '\r' && /\p{C}/u.test(c);
const isPunctuation = (c) => /[!-/:-@[-`{-~]/.test(c) || /\p{P}/u.test(c);

/** BERT normalization for an uncased model: clean, space out CJK, strip accents, lowercase one character at a time. */
export function normalize(text) {
  let out = '';
  for (const c of text) {
    const code = c.codePointAt(0);
    if (code === 0 || code === 0xFFFD || isControl(c)) continue;
    if (isSpace(c)) out += ' ';
    else out += isCjk(code) ? ` ${c} ` : c;
  }
  return Array.from(out.normalize('NFD').replace(/\p{Mn}/gu, ''), (c) => c.toLowerCase()).join('');
}

function pieces(word, vocab, unknown, ids) {
  const chars = Array.from(word);
  if (chars.length > WORD_LIMIT) return void ids.push(unknown);
  const found = [];
  for (let start = 0; start < chars.length;) {
    let end = chars.length;
    let id;
    for (; end > start; end -= 1) {
      id = vocab.get((start ? '##' : '') + chars.slice(start, end).join(''));
      if (id !== undefined) break;
    }
    if (id === undefined) return void ids.push(unknown);
    found.push(id);
    start = end;
  }
  ids.push(...found);
}

export function loadVocab(text) {
  return new Map(text.split('\n').filter((line, index, all) => line || index < all.length - 1).map((token, index) => [token, index]));
}

/** Token ids with [CLS] and [SEP], cut to `limit`, matching the tokenizer the models were trained with. */
export function tokenize(text, vocab, limit = 256) {
  const unknown = vocab.get('[UNK]');
  const room = limit - 2;
  const ids = [];
  let word = '';
  const flush = () => {
    if (word) pieces(word, vocab, unknown, ids);
    word = '';
  };
  for (const c of normalize(text)) {
    if (ids.length >= room) break;
    if (isSpace(c)) flush();
    else if (isPunctuation(c)) {
      flush();
      pieces(c, vocab, unknown, ids);
    } else word += c;
  }
  flush();
  return [vocab.get('[CLS]'), ...ids.slice(0, room), vocab.get('[SEP]')];
}

/** What the message's own headers declare. Only the part above the body counts. */
export function isBulk(text) {
  const marker = BODY_MARKER.exec(text);
  return BULK_HEADER.test(marker ? text.slice(0, marker.index) : text);
}

/** The two facts the models read beside the text. */
export const factInputs = ({ senderKnown, bulk }) => [senderKnown === true ? 1 : 0, bulk ? 1 : 0];

function lane(probabilities, bulk, acceptsPosts, rules) {
  const values = Object.fromEntries(CATEGORIES.map((name, index) => [name, probabilities[index]]));
  if (bulk) {
    values.needs_reply = 0;
    const total = CATEGORIES.reduce((sum, name) => sum + values[name], 0);
    if (total > 0) for (const name of CATEGORIES) values[name] /= total;
    else values.updates = 1;
  }
  const chosen = CATEGORIES.reduce((best, name) => (values[name] > values[best] ? name : best));
  if (chosen === 'promotional' && acceptsPosts) return 'updates';
  if (chosen === 'sales' && values.needs_reply >= rules.salesReplyFloor) return 'needs_reply';
  return chosen;
}

const top = (values) => values.reduce((best, value, index) => (value > values[best] ? index : best), 0);

/**
 * The student's answers and the facts the model cannot be argued out of, as one decision.
 * answers: { category: [5], urgency: [5], personal: number }. facts: { bulk, acceptsPosts, senderKnown }.
 */
export function decide(answers, facts, rules) {
  const { bulk, acceptsPosts, senderKnown } = facts;
  let category = lane(answers.category, bulk, acceptsPosts, rules);
  // Marketing never accepts posts, so a discussion list is the one kind of bulk mail that can wait on this person.
  const listReply = bulk && acceptsPosts && answers.category[0] >= rules.listReplyFloor;
  if (listReply) category = 'needs_reply';
  let urgency = 1;
  if (category === 'needs_reply') {
    urgency = top(answers.urgency) + 1;
    if (senderKnown === false) urgency = Math.min(urgency, rules.unknownSenderMaxUrgency);
  }
  return { category, urgency, personal: answers.personal >= rules.personalFloor };
}

/** A message can raise its score but cannot make itself trusted. */
export function urgentVerdict(score, { senderKnown, bulk }, rules) {
  if (score < rules.actionThreshold) return 'clear';
  return senderKnown === true && !bulk ? 'urgent' : 'unverified';
}

	// The Wildroot Mail engine: the tokenizer and rules around the student models. No dependencies.
	// The host runs the ONNX files. This file turns mail into model inputs and model answers into a decision.

	export const CATEGORIES = ['needs_reply', 'updates', 'promotional', 'sales', 'spam'];
	const BODY_MARKER = /(?:^\|\n)Body:\s*\n/;
	const BULK_HEADER = /(?:^\|\n)(?:Has unsubscribe header: yes\|Mailing list: [^\n]+\|Declared automated or bulk: yes)(?=\n\|$)/;
	const CJK = [[0x4E00, 0x9FFF], [0x3400, 0x4DBF], [0x20000, 0x2A6DF], [0x2A700, 0x2B73F], [0x2B740, 0x2B81F],
	[0x2B920, 0x2CEAF], [0xF900, 0xFAFF], [0x2F800, 0x2FA1F]];
	const WORD_LIMIT = 100;

	const isCjk = (code) => CJK.some(([low, high]) => code >= low && code <= high);
	const isSpace = (c) => c === '\t' \|\| c === '\n' \|\| c === '\r' \|\| /\p{White_Space}/u.test(c);
	const isControl = (c) => c !== '\t' && c !== '\n' && c !== '\r' && /\p{C}/u.test(c);
	const isPunctuation = (c) => /[!-/:-@[-`{-~]/.test(c) \|\| /\p{P}/u.test(c);

	/** BERT normalization for an uncased model: clean, space out CJK, strip accents, lowercase one character at a time. */
	export function normalize(text) {
	let out = '';
	for (const c of text) {
	const code = c.codePointAt(0);
	if (code === 0 \|\| code === 0xFFFD \|\| isControl(c)) continue;
	if (isSpace(c)) out += ' ';
	else out += isCjk(code) ? ` ${c} ` : c;
	}
	return Array.from(out.normalize('NFD').replace(/\p{Mn}/gu, ''), (c) => c.toLowerCase()).join('');
	}

	function pieces(word, vocab, unknown, ids) {
	const chars = Array.from(word);
	if (chars.length > WORD_LIMIT) return void ids.push(unknown);
	const found = [];
	for (let start = 0; start < chars.length;) {
	let end = chars.length;
	let id;
	for (; end > start; end -= 1) {
	id = vocab.get((start ? '##' : '') + chars.slice(start, end).join(''));
	if (id !== undefined) break;
	}
	if (id === undefined) return void ids.push(unknown);
	found.push(id);
	start = end;
	}
	ids.push(...found);
	}

	export function loadVocab(text) {
	return new Map(text.split('\n').filter((line, index, all) => line \|\| index < all.length - 1).map((token, index) => [token, index]));
	}

	/** Token ids with [CLS] and [SEP], cut to `limit`, matching the tokenizer the models were trained with. */
	export function tokenize(text, vocab, limit = 256) {
	const unknown = vocab.get('[UNK]');
	const room = limit - 2;
	const ids = [];
	let word = '';
	const flush = () => {
	if (word) pieces(word, vocab, unknown, ids);
	word = '';
	};
	for (const c of normalize(text)) {
	if (ids.length >= room) break;
	if (isSpace(c)) flush();
	else if (isPunctuation(c)) {
	flush();
	pieces(c, vocab, unknown, ids);
	} else word += c;
	}
	flush();
	return [vocab.get('[CLS]'), ...ids.slice(0, room), vocab.get('[SEP]')];
	}

	/** What the message's own headers declare. Only the part above the body counts. */
	export function isBulk(text) {
	const marker = BODY_MARKER.exec(text);
	return BULK_HEADER.test(marker ? text.slice(0, marker.index) : text);
	}

	/** The two facts the models read beside the text. */
	export const factInputs = ({ senderKnown, bulk }) => [senderKnown === true ? 1 : 0, bulk ? 1 : 0];

	function lane(probabilities, bulk, acceptsPosts, rules) {
	const values = Object.fromEntries(CATEGORIES.map((name, index) => [name, probabilities[index]]));
	if (bulk) {
	values.needs_reply = 0;
	const total = CATEGORIES.reduce((sum, name) => sum + values[name], 0);
	if (total > 0) for (const name of CATEGORIES) values[name] /= total;
	else values.updates = 1;
	}
	const chosen = CATEGORIES.reduce((best, name) => (values[name] > values[best] ? name : best));
	if (chosen === 'promotional' && acceptsPosts) return 'updates';
	if (chosen === 'sales' && values.needs_reply >= rules.salesReplyFloor) return 'needs_reply';
	return chosen;
	}

	const top = (values) => values.reduce((best, value, index) => (value > values[best] ? index : best), 0);

	/**
	* The student's answers and the facts the model cannot be argued out of, as one decision.
	* answers: { category: [5], urgency: [5], personal: number }. facts: { bulk, acceptsPosts, senderKnown }.
	*/
	export function decide(answers, facts, rules) {
	const { bulk, acceptsPosts, senderKnown } = facts;
	let category = lane(answers.category, bulk, acceptsPosts, rules);
	// Marketing never accepts posts, so a discussion list is the one kind of bulk mail that can wait on this person.
	const listReply = bulk && acceptsPosts && answers.category[0] >= rules.listReplyFloor;
	if (listReply) category = 'needs_reply';
	let urgency = 1;
	if (category === 'needs_reply') {
	urgency = top(answers.urgency) + 1;
	if (senderKnown === false) urgency = Math.min(urgency, rules.unknownSenderMaxUrgency);
	}
	return { category, urgency, personal: answers.personal >= rules.personalFloor };
	}

	/** A message can raise its score but cannot make itself trusted. */
	export function urgentVerdict(score, { senderKnown, bulk }, rules) {
	if (score < rules.actionThreshold) return 'clear';
	return senderKnown === true && !bulk ? 'urgent' : 'unverified';
	}