The patch introducing loading from file descriptor was rewritten to reuse code upstream. The new version honors the use_mmap when the old one did not. This makes sure we're not changing the behavior of Link Preview (only consumer) which used to request it while it was silently ignored. Old behavior: pass use_mmap=true, ignored -> no mmap. New behavior: pass use_mmap=false, honored -> no mmap. Differential Revision: https://phabricator.services.mozilla.com/D308367
634 lines
18 KiB
JavaScript
634 lines
18 KiB
JavaScript
/**
|
|
* 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
|
|
* file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
*/
|
|
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 =
|
|
"You're an AI assistant for text re-writing and summarization. Rewrite the input text focusing on the main key point in at most three very short sentences.";
|
|
|
|
// 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 FEATURE_ID = "link-preview";
|
|
|
|
const lazy = {};
|
|
ChromeUtils.defineESModuleGetters(lazy, {
|
|
createEngine: "chrome://global/content/ml/EngineProcess.sys.mjs",
|
|
FEATURES: "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"
|
|
);
|
|
XPCOMUtils.defineLazyPreferenceGetter(
|
|
lazy,
|
|
"preUserPrompt",
|
|
"browser.ml.linkPreview.preUserPrompt",
|
|
""
|
|
);
|
|
XPCOMUtils.defineLazyPreferenceGetter(
|
|
lazy,
|
|
"postUserPrompt",
|
|
"browser.ml.linkPreview.postUserPrompt",
|
|
""
|
|
);
|
|
|
|
XPCOMUtils.defineLazyPreferenceGetter(
|
|
lazy,
|
|
"penalizedTokens",
|
|
"browser.ml.linkPreview.penalizedTokens",
|
|
// default (when PREF_INVALID)
|
|
// Tokens with newlines for the default link preview model, based on the model's vocab: https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct/raw/main/vocab.json
|
|
JSON.stringify([
|
|
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,
|
|
]),
|
|
null, // no onUpdate callback
|
|
rawValue => {
|
|
if (!rawValue) {
|
|
return [];
|
|
}
|
|
return JSON.parse(rawValue);
|
|
}
|
|
);
|
|
|
|
XPCOMUtils.defineLazyPreferenceGetter(
|
|
lazy,
|
|
"minWordsPerOutputSentences",
|
|
"browser.ml.linkPreview.minWordsPerOutputSentences",
|
|
0
|
|
);
|
|
|
|
// End of generation tokens.
|
|
XPCOMUtils.defineLazyPreferenceGetter(
|
|
lazy,
|
|
"stopTokens",
|
|
"browser.ml.linkPreview.stopTokens",
|
|
// default (when PREF_INVALID)
|
|
JSON.stringify([END_OF_TEXT_TOKEN, BOS_TOKEN, EOS_TOKEN]),
|
|
null, // no onUpdate callback
|
|
rawValue => {
|
|
if (!rawValue) {
|
|
return [];
|
|
}
|
|
return JSON.parse(rawValue);
|
|
}
|
|
);
|
|
|
|
export const LinkPreviewModel = {
|
|
/**
|
|
* Manager for the block list. If null, no block list is applied.
|
|
*
|
|
* @type {BlockListManager}
|
|
*/
|
|
blockListManager: null,
|
|
|
|
get id() {
|
|
return FEATURE_ID;
|
|
},
|
|
|
|
get engineId() {
|
|
return lazy.FEATURES[this.id].engineId;
|
|
},
|
|
|
|
/**
|
|
* Blocked token list
|
|
*
|
|
* @returns {Array<number>} block token list
|
|
*/
|
|
getBlockTokenList() {
|
|
return lazy.penalizedTokens;
|
|
},
|
|
/**
|
|
* 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
|
|
) {
|
|
// Filter out emoji characters. The `u` flag is for unicode and `g` for global.
|
|
// Use `Emoji_Presentation` to avoid removing numbers and other symbols.
|
|
const textWithoutEmoji = text.replace(/\p{Emoji_Presentation}/gu, "");
|
|
return (
|
|
this.getSentences(textWithoutEmoji)
|
|
.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.
|
|
* @param {AbortSignal} abortSignal - The signal to abort the download.
|
|
* @returns {Promise<MLEngine>} - A promise that resolves to the ML engine instance.
|
|
*/
|
|
async createEngine(options, notificationsCallback = null, abortSignal) {
|
|
return lazy.createEngine(options, notificationsCallback, abortSignal);
|
|
},
|
|
|
|
/**
|
|
* Generate summary text using AI.
|
|
*
|
|
* @param {string} inputText
|
|
* @param {object} callbacks for progress and error
|
|
* @param {AbortSignal} callbacks.abortSignal - The signal to abort the download.
|
|
* @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, abortSignal } = {}
|
|
) {
|
|
// 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: "best-llama",
|
|
engineId: this.engineId,
|
|
featureId: this.id,
|
|
kvCacheDtype: "q8_0",
|
|
modelFile: "smollm2-360m-instruct-q8_0.gguf",
|
|
modelHubRootUrl: "https://model-hub.mozilla.org",
|
|
modelId: "HuggingFaceTB/SmolLM2-360M-Instruct-GGUF",
|
|
modelRevision: "main",
|
|
numBatch: numContext,
|
|
numContext,
|
|
numUbatch: numContext,
|
|
taskName: "wllama-text-generation",
|
|
timeoutMS: -1,
|
|
useMlock: false,
|
|
useMmap: false,
|
|
...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)
|
|
);
|
|
}
|
|
},
|
|
abortSignal
|
|
);
|
|
|
|
const postProcessor = await SentencePostProcessor.initialize();
|
|
const blockedTokens = this.getBlockTokenList();
|
|
for await (const val of engine.runWithGenerator({
|
|
nPredict,
|
|
stopTokens: lazy.stopTokens,
|
|
logit_bias_toks: blockedTokens,
|
|
logit_bias_vals: Array(blockedTokens.length).fill(-Infinity),
|
|
prompt: [
|
|
{ role: "system", content: systemPrompt },
|
|
{
|
|
role: "user",
|
|
content: lazy.preUserPrompt + processedInput + lazy.postUserPrompt,
|
|
},
|
|
],
|
|
})) {
|
|
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("");
|
|
// simple way to get number of word ignoring non-whitespaces chatacters
|
|
const isValidSentence =
|
|
sentences[0].trim().split(/\p{White_Space}+/u).length >=
|
|
lazy.minWordsPerOutputSentences;
|
|
|
|
if (isValidSentence) {
|
|
this.currentNumSentences += 1;
|
|
}
|
|
|
|
if (this.currentNumSentences == this.maxNumOutputSentences) {
|
|
this.currentText = "";
|
|
abort = true;
|
|
}
|
|
if (isValidSentence) {
|
|
sentence = sentences[0];
|
|
}
|
|
|
|
// If the sentence contains a block word, abort
|
|
if (
|
|
this.blockListManager &&
|
|
this.blockListManager.matchAtWordBoundary({
|
|
// Blocklist is always lowercase
|
|
text: sentence.toLowerCase(),
|
|
})
|
|
) {
|
|
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;
|
|
}
|
|
}
|