Source code
Revision control
Copy as Markdown
Other Tools
/**
* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
*/
import { XPCOMUtils } from "resource://gre/modules/XPCOMUtils.sys.mjs";
// On average, each token represents about 4 characters. A factor of 3.5 is used
// instead of 4 to account for edge cases.
const CHARACTERS_PER_TOKEN = 3.5;
// On average, one token corresponds to approximately 4 characters, meaning 0.25
// times the character count would suffice under normal conditions. To ensure
// robustness and handle edge cases, we use a more conservative factor of 0.69.
const CONTEXT_SIZE_MULTIPLIER = 0.69;
const DEFAULT_INPUT_SENTENCES = 6;
const MIN_SENTENCE_LENGTH = 14;
const MIN_WORD_COUNT = 5;
const DEFAULT_INPUT_PROMPT =
"Provide a concise, objective summary of the input text in up to three sentences, focusing on key actions and intentions without using second or third person pronouns.";
// All tokens taken from the model's vocabulary at https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct/raw/main/vocab.json
// Token id for end of text
const END_OF_TEXT_TOKEN = 0;
// Token id for beginning of sequence
const BOS_TOKEN = 1;
// Token id for end of sequence
const EOS_TOKEN = 2;
const lazy = {};
ChromeUtils.defineESModuleGetters(lazy, {
createEngine: "chrome://global/content/ml/EngineProcess.sys.mjs",
Progress: "chrome://global/content/ml/Utils.sys.mjs",
BlockListManager: "chrome://global/content/ml/Utils.sys.mjs",
RemoteSettingsManager: "chrome://global/content/ml/Utils.sys.mjs",
});
XPCOMUtils.defineLazyPreferenceGetter(
lazy,
"config",
"browser.ml.linkPreview.config",
"{}"
);
XPCOMUtils.defineLazyPreferenceGetter(
lazy,
"inputSentences",
"browser.ml.linkPreview.inputSentences"
);
XPCOMUtils.defineLazyPreferenceGetter(
lazy,
"outputSentences",
"browser.ml.linkPreview.outputSentences"
);
XPCOMUtils.defineLazyPreferenceGetter(
lazy,
"prompt",
"browser.ml.linkPreview.prompt"
);
XPCOMUtils.defineLazyPreferenceGetter(
lazy,
"blockListEnabled",
"browser.ml.linkPreview.blockListEnabled"
);
export const LinkPreviewModel = {
/**
* Manager for the block list. If null, no block list is applied.
*
* @type {BlockListManager}
*/
blockListManager: null,
/**
* Blocked token list
*
* @returns {Array<number>} block token list
*/
getBlockTokenList() {
// Tokens with newlines for the link preview model, based on the model's vocab: https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct/raw/main/vocab.json
const tokensWithNewLines = [
198, 448, 466, 472, 629, 945, 1004, 1047, 1116, 1410, 1927, 2367, 2738,
2830, 2953, 3136, 3299, 3337, 3354, 3558, 3717, 3805, 3914, 4602, 4767,
5952, 7116, 7209, 7338, 7396, 8301, 8500, 8821, 8866, 9198, 9225, 9343,
9694, 10459, 11181, 11259, 11539, 11813, 12350, 13002, 13272, 13280,
13596, 13617, 13809, 14436, 14446, 15111, 15182, 15290, 15537, 16140,
16299, 16390, 16506, 16871, 16980, 16997, 18682, 18850, 18864, 19014,
19145, 19993, 20098, 20370, 20793, 21193, 21377, 21941, 22342, 22369,
23004, 23386, 23499, 23799, 24112, 24205, 25457, 25576, 26675, 26886,
26925, 27536, 27924, 28577, 29306, 29866, 30314, 30544, 30799, 31464,
32057, 32315, 32829, 34344, 34356, 35163, 35988, 36176, 36286, 36328,
36489, 36496, 36804, 37468, 38028, 38031, 39014, 39843, 39892, 40677,
40944, 42057, 42617, 43784, 43902, 44064, 46778, 47213, 47647, 48259,
48279, 48818,
];
return tokensWithNewLines;
},
/**
* Extracts sentences from a given text.
*
* @param {string} text text to process
* @returns {Array<string>} sentences
*/
getSentences(text) {
const abbreviations = [
"Mr.",
"Mrs.",
"Ms.",
"Dr.",
"Prof.",
"Inc.",
"Ltd.",
"Jr.",
"Sr.",
"St.",
"e.g.",
"i.e.",
"U.S.A",
"D.C.",
"U.K.",
"etc.",
"a.m.",
"p.m.",
"D.",
"Mass.",
"Sen.",
"Rep.",
"No.",
"Fig.",
"vs.",
"Mx.",
"Ph.D.",
"M.D.",
"D.D.S.",
"B.A.",
"M.A.",
"LL.B.",
"LL.M.",
"J.D.",
"D.O.",
"D.V.M.",
"Psy.D.",
"Ed.D.",
"Eng.",
"Co.",
"Corp.",
"Mt.",
"Ft.",
"U.S.",
"U.S.A.",
"E.U.",
"et al.",
"Nos.",
"pp.",
"Vol.",
"Rev.",
"Gen.",
"Lt.",
"Col.",
"Maj.",
"Capt.",
"Sgt.",
"Cpl.",
"Pvt.",
"Adm.",
"Cmdr.",
"Ave.",
"Blvd.",
"Rd.",
"Ln.",
"Jan.",
"Feb.",
"Mar.",
"Apr.",
"May.",
"Jun.",
"Jul.",
"Aug.",
"Sep.",
"Sept.",
"Oct.",
"Nov.",
"Dec.",
"Mon.",
"Tue.",
"Tues.",
"Wed.",
"Thu.",
"Thur.",
"Thurs.",
"Fri.",
"Sat.",
"Sun.",
"Dept.",
"Univ.",
"Est.",
"Calif.",
"Fla.",
"N.Y.",
"Conn.",
"Va.",
"Ill.",
"Assoc.",
"Bros.",
"Dist.",
"Msgr.",
"S.P.",
"P.S.",
"U.S.S.R.",
"Mlle.",
"Mme.",
"Hon.",
"Messrs.",
"Mmes.",
"v.",
"vs.",
];
// Replace periods in abbreviations with a placeholder.
let modifiedText = text;
const placeholder = "∯";
abbreviations.forEach(abbrev => {
const escapedAbbrev = abbrev
.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")
.replace(/\\\./g, "\\.");
const regex = new RegExp(escapedAbbrev, "g");
const abbrevWithPlaceholder = abbrev.replace(/\./g, placeholder);
modifiedText = modifiedText.replace(regex, abbrevWithPlaceholder);
});
const segmenter = new Intl.Segmenter("en", {
granularity: "sentence",
});
const segments = segmenter.segment(modifiedText);
let sentences = Array.from(segments, segment => segment.segment);
// Restore the periods in abbreviations.
return sentences.map(sentence =>
sentence.replace(new RegExp(placeholder, "g"), ".")
);
},
/**
* Clean up text for text generation AI.
*
* @param {string} text to process
* @param {number} maxNumSentences - Max number of sentences to return.
* @returns {string} cleaned up text
*/
preprocessText(
text,
maxNumSentences = lazy.inputSentences ?? DEFAULT_INPUT_SENTENCES
) {
return (
this.getSentences(text)
.map(s =>
// trim and replace consecutive blank by a single one.
s.trim().replace(
/(\s*\n\s*)|\s{2,}/g,
// (\s*\n\s*) -> Matches a newline (`\n`) surrounded by optional whitespace.
// \s{2,} -> Matches two or more consecutive spaces.
// g -> Global flag to replace all occurrences in the string.
(_, newline) => (newline ? "\n" : " ")
// Callback function:
// `_` -> First argument (full match) is ignored.
// `newline` -> If the first capturing group (\s*\n\s*) matched, `newline` is truthy.
// If `newline` exists, it replaces the match with a single newline ("\n").
// Otherwise, it replaces the match (extra spaces) with a single space (" ").
)
)
// Remove sentences that are too short without punctuation.
.filter(
s =>
s.length >= MIN_SENTENCE_LENGTH &&
s.split(" ").length >= MIN_WORD_COUNT &&
/\p{P}$/u.test(s)
)
.slice(0, maxNumSentences)
.join(" ")
);
},
/**
* Creates a new ML engine instance with the provided options for link preview.
*
* @param {object} options - Configuration options for the ML engine.
* @param {?function(ProgressAndStatusCallbackParams):void} notificationsCallback A function to call to indicate notifications.
* @returns {Promise<MLEngine>} - A promise that resolves to the ML engine instance.
*/
async createEngine(options, notificationsCallback = null) {
return lazy.createEngine(options, notificationsCallback);
},
/**
* Generate summary text using AI.
*
* @param {string} inputText
* @param {object} callbacks for progress and error
* @param {Function} callbacks.onDownload optional for download active
* @param {Function} callbacks.onText optional for text chunks
* @param {Function} callbacks.onError optional for error
*/
async generateTextAI(inputText, { onDownload, onText, onError } = {}) {
// Get updated options from remote settings. No failure if no record exists
const remoteRequestRecord = await lazy.RemoteSettingsManager.getRemoteData({
collectionName: "ml-inference-request-options",
filters: { featureId: "link-preview" },
majorVersion: 1,
}).catch(() => {
console.error(
"Error retrieving request options from remote settings, will use default options."
);
return { options: "{}" };
});
let remoteRequestOptions = {};
try {
remoteRequestOptions = remoteRequestRecord?.options
? JSON.parse(remoteRequestRecord.options)
: {};
} catch (error) {
console.error(
"Error parsing the remote settings request options, will use default options.",
error
);
}
// TODO: Unit test that order of preference is correctly respected.
const processedInput = this.preprocessText(
inputText,
lazy.inputSentences ??
remoteRequestOptions?.inputSentences ??
DEFAULT_INPUT_SENTENCES
);
// Asssume generated text is approximately the same length as the input.
const nPredict = Math.ceil(processedInput.length / CHARACTERS_PER_TOKEN);
const systemPrompt =
lazy.prompt ?? remoteRequestOptions?.systemPrompt ?? DEFAULT_INPUT_PROMPT;
// Estimate an upper bound for the required number of tokens. This estimate
// must be large enough to include prompt tokens, input tokens, and
// generated tokens.
const numContext =
Math.ceil(
(processedInput.length + systemPrompt.length) * CONTEXT_SIZE_MULTIPLIER
) + nPredict;
let engine;
try {
engine = await this.createEngine(
{
backend: "wllama",
engineId: "wllamapreview",
kvCacheDtype: "q8_0",
modelFile: "smollm2-360m-instruct-q8_0.gguf",
modelHubUrlTemplate: "{model}/{revision}",
modelId: "HuggingFaceTB/SmolLM2-360M-Instruct-GGUF",
modelRevision: "main",
numBatch: numContext,
numContext,
numUbatch: numContext,
runtimeFilename: "wllama.wasm",
taskName: "wllama-text-generation",
timeoutMS: -1,
useMlock: false,
useMmap: true,
...JSON.parse(lazy.config),
},
data => {
if (data.type == lazy.Progress.ProgressType.DOWNLOAD) {
onDownload?.(
data.statusText != lazy.Progress.ProgressStatusText.DONE,
Math.round((100 * data.totalLoaded) / data.total)
);
}
}
);
const postProcessor = await SentencePostProcessor.initialize();
const blockedTokens = this.getBlockTokenList();
for await (const val of engine.runWithGenerator({
nPredict,
stopTokens: [END_OF_TEXT_TOKEN, BOS_TOKEN, EOS_TOKEN],
logit_bias_toks: blockedTokens,
logit_bias_vals: Array(blockedTokens.length).fill(-Infinity),
prompt: [
{ role: "system", content: systemPrompt },
{ role: "user", content: processedInput },
],
})) {
const { sentence, abort } = postProcessor.put(val.text);
if (sentence) {
onText?.(sentence);
} else if (!val.text) {
const remaining = postProcessor.flush();
if (remaining) {
onText?.(remaining);
}
}
if (abort) {
break;
}
}
} catch (error) {
onError?.(error);
} finally {
await engine?.terminate();
}
},
};
/**
* A class for processing streaming text to detect and extract complete
* sentences. It buffers incoming text and periodically checks for new sentences
* based on punctuation and character count limits.
*
* This class is useful for incremental sentence processing in NLP tasks.
*/
export class SentencePostProcessor {
/**
* The maximum number of sentences to output before truncating the buffer.
* Use -1 for unlimited.
*
* @type {number}
*/
maxNumOutputSentences = -1;
/**
* Stores the current text being processed.
*
* @type {string}
*/
currentText = "";
/**
* Tracks the number of sentences processed so far.
*
* @type {number}
*/
currentNumSentences = 0;
/**
* Manager for the block list. If null, no block list is applied.
*
* @type {BlockListManager}
*/
blockListManager = null;
/**
* Create an instance of the sentence postprocessor.
*
* @param {object} config - Configuration object.
* @param {number} config.maxNumOutputSentences - The maximum number of sentences to
* output before truncating the buffer.
* @param {BlockListManager | null} config.blockListManager - Manager for the block list
*/
constructor({
maxNumOutputSentences = lazy.outputSentences,
blockListManager,
} = {}) {
this.maxNumOutputSentences = maxNumOutputSentences;
this.blockListManager = blockListManager;
}
/**
* @param {object} config - Configuration object.
* @param {number} config.maxNumOutputSentences - The maximum number of sentences to
* output before truncating the buffer.
* @param {boolean} config.blockListEnabled - Wether to enable block list. If enabled, we
* don't return the sentence that has a blocked word along with any sentences coming after.
* @returns {SentencePostProcessor} - An instance of SentencePostProcessor
*/
static async initialize({
maxNumOutputSentences = lazy.outputSentences,
blockListEnabled = lazy.blockListEnabled,
} = {}) {
if (!blockListEnabled) {
LinkPreviewModel.blockListManager = null;
} else if (!LinkPreviewModel.blockListManager) {
LinkPreviewModel.blockListManager =
await lazy.BlockListManager.initializeFromRemoteSettings({
blockListName: "link-preview-test-en",
language: "en",
fallbackToDefault: true,
majorVersion: 1,
});
}
return new SentencePostProcessor({
maxNumOutputSentences,
blockListManager: LinkPreviewModel.blockListManager,
});
}
/**
* Processes incoming text, checking if a full sentence has been completed. If
* a full sentence is detected, it returns the first complete sentence.
* Otherwise, it returns an empty string.
*
* @param {string} text to process
* @returns {{ text: string, abort: boolean }} An object containing:
* - `{string} sentence`: The first complete sentence if available, otherwise an empty string.
* - `{boolean} abort`: `true` if generation should be aborted early, `false` otherwise.
*/
put(text) {
if (this.currentNumSentences == this.maxNumOutputSentences) {
return { sentence: "", abort: true };
}
this.currentText += text;
// We need to ensure that the current sentence is complete and the next
// has started before reporting that a sentence is ready.
const sentences = LinkPreviewModel.getSentences(this.currentText);
let sentence = "";
let abort = false;
if (sentences.length >= 2) {
this.currentText = sentences.slice(1).join("");
this.currentNumSentences += 1;
if (this.currentNumSentences == this.maxNumOutputSentences) {
this.currentText = "";
abort = true;
}
sentence = sentences[0];
// If the sentence contains a block word, abort
if (
this.blockListManager &&
this.blockListManager.matchAtWordBoundary({ text: sentence })
) {
sentence = "";
abort = true;
this.currentNumSentences = this.maxNumOutputSentences;
}
}
return { sentence, abort };
}
/**
* Flushes the remaining text buffer. This ensures that any last remaining
* sentence is returned.
*
* @returns {string} remaining text that hasn't been processed yet
*/
flush() {
return this.currentText;
}
}