1278 lines
35 KiB
JavaScript
1278 lines
35 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/. */
|
|
|
|
// @ts-nocheck - TODO - Remove this to type check this file.
|
|
|
|
/**
|
|
* @typedef {import("../actors/MLEngineParent.sys.mjs").MLEngineParent} MLEngineParent
|
|
* @typedef {import("../content/Utils.sys.mjs").ProgressAndStatusCallbackParams} ProgressAndStatusCallbackParams
|
|
*/
|
|
|
|
/**
|
|
* @constant
|
|
* @type {string}
|
|
* @default
|
|
* @description The default engine identifier used when no specific engine ID is provided.
|
|
*/
|
|
export const DEFAULT_ENGINE_ID = "default-engine";
|
|
|
|
/**
|
|
* Supported backends.
|
|
*/
|
|
export const BACKENDS = Object.freeze({
|
|
onnx: "onnx",
|
|
wllama: "wllama",
|
|
onnxNative: "onnx-native",
|
|
llamaCpp: "llama.cpp",
|
|
bestLlama: "best-llama",
|
|
bestOnnx: "best-onnx",
|
|
openai: "openai",
|
|
staticEmbeddings: "static-embeddings",
|
|
});
|
|
|
|
/**
|
|
* @constant
|
|
* @type {Array<string>}
|
|
* @description Backends using WASM.
|
|
*/
|
|
export const WASM_BACKENDS = [BACKENDS.onnx, BACKENDS.wllama];
|
|
|
|
/**
|
|
* @constant
|
|
* @type {string}
|
|
* @description Matches filenames with subdirectories, starting with alphanumeric or underscore,
|
|
and optionally ending with a dot followed by a 2-9 letter extension.
|
|
|
|
^ $ Start and end of string
|
|
(?:\/)? Optional leading slash (for absolute paths or root directory)
|
|
(?!\/) Negative lookahead for not starting with a slash
|
|
[A-Za-z0-9-_]+ First directory or filename
|
|
(?: Begin non-capturing group for additional directories or file
|
|
\/ Directory separator
|
|
[A-Za-z0-9-_.]+ Directory or file name
|
|
)* Zero or more times
|
|
(?:[.][A-Za-z_]{2,9})? Optional non-capturing group for file extension
|
|
*/
|
|
export const FILE_REGEX =
|
|
/^(?:\/)?(?!\/)[A-Za-z0-9-_.]+(?:\/[A-Za-z0-9-_.]+)*(?:[.][A-Za-z_]{2,9})?$/;
|
|
|
|
/**
|
|
* @constant
|
|
* @type {{ [key: string]: { modelId: string, dtype: string } }}
|
|
* @description Supported tasks with their default model identifiers.
|
|
*/
|
|
export const DEFAULT_MODELS = Object.freeze({
|
|
"test-echo": { modelId: "test-echo", dtype: "q8" },
|
|
"text-classification": {
|
|
modelId: "Xenova/distilbert-base-uncased-finetuned-sst-2-english",
|
|
dtype: "q8",
|
|
},
|
|
"token-classification": {
|
|
modelId: "Xenova/bert-base-multilingual-cased-ner-hrl",
|
|
dtype: "q8",
|
|
},
|
|
"question-answering": {
|
|
modelId: "Xenova/distilbert-base-cased-distilled-squad",
|
|
dtype: "q8",
|
|
},
|
|
"fill-mask": { modelId: "Xenova/bert-base-uncased", dtype: "q8" },
|
|
summarization: { modelId: "Xenova/distilbart-cnn-6-6", dtype: "q8" },
|
|
translation: { modelId: "Xenova/t5-small", dtype: "q8" },
|
|
"text2text-generation": { modelId: "Xenova/flan-t5-small", dtype: "q8" },
|
|
"text-generation": { modelId: "onnx-community/gpt2-ONNX", dtype: "q8" },
|
|
"zero-shot-classification": {
|
|
modelId: "Xenova/distilbert-base-uncased-mnli",
|
|
dtype: "q8",
|
|
},
|
|
"image-to-text": { modelId: "Mozilla/distilvit", dtype: "q8" },
|
|
"image-classification": {
|
|
modelId: "Xenova/vit-base-patch16-224",
|
|
dtype: "q8",
|
|
},
|
|
"image-segmentation": {
|
|
modelId: "Xenova/detr-resnet-50-panoptic",
|
|
dtype: "q8",
|
|
},
|
|
"zero-shot-image-classification": {
|
|
modelId: "Xenova/clip-vit-base-patch32",
|
|
dtype: "q8",
|
|
},
|
|
"object-detection": { modelId: "Xenova/detr-resnet-50", dtype: "q8" },
|
|
"zero-shot-object-detection": {
|
|
modelId: "Xenova/owlvit-base-patch32",
|
|
dtype: "q8",
|
|
},
|
|
"document-question-answering": {
|
|
modelId: "Xenova/donut-base-finetuned-docvqa",
|
|
dtype: "q8",
|
|
},
|
|
"image-to-image": {
|
|
modelId: "Xenova/swin2SR-classical-sr-x2-64",
|
|
dtype: "q8",
|
|
},
|
|
"depth-estimation": { modelId: "Xenova/dpt-large", dtype: "q8" },
|
|
"feature-extraction": {
|
|
modelId: "Xenova/all-MiniLM-L6-v2",
|
|
dtype: "q8",
|
|
},
|
|
"image-feature-extraction": {
|
|
modelId: "Xenova/vit-base-patch16-224-in21k",
|
|
dtype: "q8",
|
|
},
|
|
"text-to-speech": {
|
|
modelId: "Xenova/speecht5_tts",
|
|
dtype: "q8",
|
|
},
|
|
});
|
|
|
|
/**
|
|
* Lists Firefox internal features.
|
|
*
|
|
* Make sure to set an explicit engine id and add it in toolkit/components/ml/metrics.yml
|
|
*
|
|
* Engine ids can be shared, there's only one inference runtime running per engine id.
|
|
* When an inference is executed on an engine, if the pipeline options are different
|
|
* from the previous call, the engine is re-created. If they are the same, the engine is reused.
|
|
*
|
|
* The only exception is web extension, as the engine id is dynamically created with the extension id.
|
|
*/
|
|
export const FEATURES = {
|
|
// see toolkit/components/pdfjs/content/PdfjsParent.sys.mjs
|
|
"pdfjs-alt-text": {
|
|
engineId: "pdfjs",
|
|
fluentId: "mlmodel-pdfjs",
|
|
},
|
|
// see browser/components/urlbar/private/MLSuggest.sys.mjs
|
|
"suggest-intent-classification": {
|
|
engineId: "ml-suggest-intent",
|
|
},
|
|
// see browser/components/urlbar/private/MLSuggest.sys.mjs
|
|
"suggest-NER": {
|
|
engineId: "ml-suggest-ner",
|
|
},
|
|
// see toolkit/components/aboutinference/content/aboutInference.js
|
|
"about-inference": {
|
|
engineId: "about-inference",
|
|
fluentId: "mlmodel-about-inference",
|
|
},
|
|
// see browser/components/tabbrowser/SmartTabGrouping.sys.mjs,
|
|
"smart-tab-embedding": {
|
|
engineId: "smart-tab-embedding-engine",
|
|
fluentId: "mlmodel-smart-tab-embedding-engine",
|
|
},
|
|
// see browser/components/tabbrowser/SmartTabGrouping.sys.mjs
|
|
"smart-tab-topic": {
|
|
engineId: "smart-tab-topic-engine",
|
|
fluentId: "mlmodel-smart-tab-topic-engine",
|
|
},
|
|
// see toolkit/components/formautofill/shared/FormAutofillML.sys.mjs
|
|
"formfill-classification": {
|
|
engineId: "formfill-classification-engine",
|
|
fluentId: "mlmodel-formfill-engine",
|
|
},
|
|
// see toolkit/components/ml/content/nlp/EmbeddingsGenerator.sys.mjs
|
|
"simple-text-embedder": {
|
|
engineId: "simple-text-embedder-engine",
|
|
},
|
|
// see browser/components/genai/LinkPreviewModel.sys.mjs
|
|
"link-preview": {
|
|
engineId: "link-preview",
|
|
fluentId: "mlmodel-link-preview",
|
|
},
|
|
// see browser/components/aiwindow/models/IntentClassifier.sys.mjs
|
|
"smart-intent": {
|
|
engineId: "smart-intent",
|
|
},
|
|
"smart-intent-en-fr": {
|
|
engineId: "smart-intent-en-fr",
|
|
},
|
|
chat: {
|
|
engineId: "smart-openai",
|
|
},
|
|
"title-generation": {
|
|
engineId: "title-generation-engine",
|
|
},
|
|
"conversation-suggestions-sidebar-starter": {
|
|
engineId: "smart-openai",
|
|
},
|
|
"conversation-suggestions-followup": {
|
|
engineId: "smart-openai",
|
|
},
|
|
"memories-initial-generation-system": {
|
|
engineId: "smart-openai-memories-generation",
|
|
},
|
|
"memories-message-classification-system": {
|
|
engineId: "smart-openai-memories-usage",
|
|
},
|
|
"llm-telemetry": {
|
|
engineId: "llm-telemetry-engine",
|
|
},
|
|
//agents
|
|
"agent-monitor": {
|
|
engineId: "agent-monitor-engine",
|
|
},
|
|
};
|
|
|
|
/**
|
|
* Custom error class for validation errors.
|
|
*
|
|
* This error is thrown when a field fails validation, providing additional context such as
|
|
* the name of the field that caused the error.
|
|
*
|
|
* @augments Error
|
|
*/
|
|
class PipelineOptionsValidationError extends Error {
|
|
/**
|
|
* Create a PipelineOptionsValidationError.
|
|
*
|
|
* @param {string} field - The name of the field that caused the validation error.
|
|
* @param {any} value - The invalid value provided for the field.
|
|
* @param {string} [tips=null] - Optional tips or suggestions for valid values.
|
|
*/
|
|
constructor(field, value, tips = null) {
|
|
const baseMessage = `Invalid value "${value}" for field "${field}".`;
|
|
const message = tips ? `${baseMessage} ${tips}` : baseMessage;
|
|
super(message);
|
|
|
|
this.name = this.constructor.name;
|
|
this.field = field;
|
|
this.value = value;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Enum for model hubs
|
|
*
|
|
* @readonly
|
|
* @enum {string}
|
|
*/
|
|
export const ModelHub = {
|
|
HUGGINGFACE: "huggingface",
|
|
MOZILLA: "mozilla",
|
|
|
|
apply(options, hub) {
|
|
switch (hub) {
|
|
case ModelHub.HUGGINGFACE:
|
|
options.modelHubRootUrl = "https://huggingface.co/";
|
|
options.modelHubUrlTemplate = "{model}/resolve/{revision}";
|
|
options.modelRevision = "main";
|
|
break;
|
|
case ModelHub.MOZILLA:
|
|
options.modelHubRootUrl = "https://model-hub.mozilla.org/";
|
|
options.modelHubUrlTemplate = "{model}/{revision}";
|
|
options.modelRevision = "main";
|
|
break;
|
|
default:
|
|
throw new Error(`Unknown model hub: ${hub}`);
|
|
}
|
|
},
|
|
};
|
|
|
|
/**
|
|
* Enum for execution priority of a task.
|
|
*
|
|
* @readonly
|
|
* @enum {string}
|
|
*/
|
|
export const ExecutionPriority = {
|
|
// The task is absolutely needed for Firefox.
|
|
HIGH: "HIGH",
|
|
// The default priority.
|
|
NORMAL: "NORMAL",
|
|
// 3rd party calls.
|
|
LOW: "LOW",
|
|
};
|
|
|
|
/**
|
|
* Enum for model quantization levels. Not all models support all values.
|
|
*
|
|
* @readonly
|
|
* @enum {string}
|
|
*/
|
|
export const QuantizationLevel = {
|
|
// Full precision 32-bit floating point (`''`)
|
|
FP32: "fp32",
|
|
// Half precision 16-bit floating point (`'_fp16'`)
|
|
FP16: "fp16",
|
|
// Floating point 8 with the exponential taking 5 bits, and the mantissa taking 2.
|
|
// This format can express a wide dynamic range of float values because of the
|
|
// extra bits in the exponential, but with the trade off of lower precision of stored
|
|
// values because of the small mantissa. The max finite values are ±57,344.
|
|
FP8_E5M2: "fp8_e5m2",
|
|
// Floating point 8 with the exponential taking 4 bits, and the mantissa taking 3.
|
|
// This format is best for values without a wide dynamic range. The higher bits
|
|
// in the mantissa retains more precision The max finite values are ±448.
|
|
FP8_E4M3: "fp8_e4m3",
|
|
// Quantized 8-bit (`'_quantized'`)
|
|
Q8: "q8",
|
|
// Integer 8-bit quantization (`'_int8'`)
|
|
INT8: "int8",
|
|
// Unsigned integer 8-bit quantization (`'_uint8'`)
|
|
UINT8: "uint8",
|
|
// Quantized 4-bit (`'_q4'`)
|
|
Q4: "q4",
|
|
// Binary/Boolean 4-bit quantization (`'_bnb4'`)
|
|
BNB4: "bnb4",
|
|
// 16-bit floating point model with 4-bit block weight quantization (`'_q4f16'`)
|
|
Q4F16: "q4f16",
|
|
};
|
|
|
|
/**
|
|
* Enum for KV cache quantization levels.
|
|
*
|
|
* @readonly
|
|
* @enum {string}
|
|
*/
|
|
export const KVCacheQuantizationLevel = {
|
|
// Quantized 8-bit with optimized storage (`'_q8_0'`) (block-based)
|
|
Q8_0: "q8_0",
|
|
// Quantized 4-bit version 0 (`'_q4_0'`) (block-based)
|
|
Q4_0: "q4_0",
|
|
// Quantized 4-bit version 1 (`'_q4_1'`) (block-based)
|
|
Q4_1: "q4_1",
|
|
// Quantized 5-bit version 1 (`'_q5_1'`) (block-based)
|
|
Q5_1: "q5_1",
|
|
// Quantized 5-bit version 0 (`'_q5_0'`) (block-based)
|
|
Q5_0: "q5_0",
|
|
// Half-precision (16-bit floating point) (`'_f16'`)
|
|
F16: "f16",
|
|
// Full precision (32-bit floating point) (`'_f32'`)
|
|
F32: "f32",
|
|
};
|
|
|
|
/**
|
|
* Enum for the device used for inference.
|
|
*
|
|
* @readonly
|
|
* @enum {string}
|
|
*/
|
|
export const InferenceDevice = {
|
|
GPU: "gpu",
|
|
WASM: "wasm",
|
|
CPU: "cpu",
|
|
};
|
|
|
|
/**
|
|
* Enum for log levels.
|
|
*
|
|
* @readonly
|
|
* @enum {string}
|
|
*/
|
|
export const LogLevel = {
|
|
TRACE: "Trace",
|
|
INFO: "Info",
|
|
DEBUG: "Debug",
|
|
WARN: "Warn",
|
|
ERROR: "Error",
|
|
CRITICAL: "Critical",
|
|
ALL: "All",
|
|
};
|
|
|
|
/**
|
|
* Allowed boolean values.
|
|
*
|
|
* @readonly
|
|
* @enum {boolean}
|
|
*/
|
|
export const AllowedBoolean = [false, true];
|
|
|
|
/**
|
|
* @import { TranslationsEngineParent } from "../../translations/actors/TranslationsEngineParent.sys.mjs"
|
|
* @import { StaticEmbeddingsOptions } from "./backends/StaticEmbeddingsPipeline.d.ts"
|
|
* @import { PURPOSES, SERVICE_TYPES } from "../../../../browser/components/aiwindow/models/Utils.sys.mjs"
|
|
*/
|
|
|
|
const PIPELINE_TEST_NAMES = ["moz-echo", "test-echo"];
|
|
|
|
/**
|
|
* This class encapsulates the options for a pipeline process.
|
|
*
|
|
* DO NOT SET A DEFAULT VALUE HERE - KEEP IT TO NULL, except engineId.
|
|
*
|
|
* Options should be set to null here, and each pipeline backend is
|
|
* responsible for setting a default value when appropriate.
|
|
*
|
|
* If there's a non null value here it might override a value defined by the
|
|
* caller.
|
|
*/
|
|
export class PipelineOptions {
|
|
/**
|
|
* The identifier for the engine to be used by the pipeline.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
engineId = DEFAULT_ENGINE_ID;
|
|
|
|
/**
|
|
* External model data file list.
|
|
*/
|
|
useExternalDataFormat = false;
|
|
|
|
/**
|
|
* The name of the feature to be used by the pipeline.
|
|
*
|
|
* This field can be used to uniquely identify an inference and
|
|
* overwrite taskName when doing lookups in Remote Settings.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
featureId = null;
|
|
|
|
/**
|
|
* The name of the task the pipeline is configured for.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
taskName = null;
|
|
|
|
/**
|
|
* The maximum amount of time in milliseconds the pipeline should wait for a response.
|
|
*
|
|
* @type {?number}
|
|
*/
|
|
timeoutMS = null;
|
|
|
|
/**
|
|
* The hub to use. When null, looks at modelHubRootUrl and modelHubUrlTemplate
|
|
*
|
|
* @type {ModelHub | null}
|
|
*/
|
|
modelHub = null;
|
|
|
|
/**
|
|
* The root URL of the model hub where models are hosted.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
modelHubRootUrl = null;
|
|
|
|
/**
|
|
* A template URL for building the full URL for the model.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
modelHubUrlTemplate = null;
|
|
|
|
/**
|
|
* The identifier for the specific model to be used by the pipeline.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
modelId = null;
|
|
|
|
/**
|
|
* The revision for the specific model to be used by the pipeline.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
modelRevision = null;
|
|
|
|
/**
|
|
* The flowId is used to track a flow of events for telemetry.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
flowId = null;
|
|
|
|
/**
|
|
* The identifier for the tokenizer associated with the model, used for pre-processing inputs.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
tokenizerId = null;
|
|
|
|
/**
|
|
* The revision for the tokenizer associated with the model, used for pre-processing inputs.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
tokenizerRevision = null;
|
|
|
|
/**
|
|
* The identifier for any processor required by the model, used for additional input processing.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
processorId = null;
|
|
|
|
/**
|
|
* The revision for any processor required by the model, used for additional input processing.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
|
|
processorRevision = null;
|
|
|
|
/**
|
|
* The log level used in the worker
|
|
*
|
|
* @type {LogLevel | null}
|
|
*/
|
|
logLevel = null;
|
|
|
|
/**
|
|
* Name of the runtime wasm file
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
runtimeFilename = null;
|
|
|
|
/**
|
|
* Device used for inference
|
|
*
|
|
* @type {InferenceDevice | null}
|
|
*/
|
|
device = null;
|
|
|
|
/**
|
|
* Model Quantization level
|
|
*
|
|
* @type {QuantizationLevel | null}
|
|
*/
|
|
dtype = null;
|
|
|
|
/**
|
|
* Number of threads to use in the pipeline
|
|
*
|
|
* @type {?number}
|
|
*/
|
|
numThreads = null;
|
|
|
|
/**
|
|
* Execution priority
|
|
*
|
|
* Defines the priority of the task
|
|
*
|
|
* @type {ExecutionPriority}
|
|
*/
|
|
executionPriority = null;
|
|
|
|
/**
|
|
* KV cache Quantization level
|
|
*
|
|
* @type {KVCacheQuantizationLevel | null}
|
|
*/
|
|
kvCacheDtype = null;
|
|
|
|
/**
|
|
* Maximum context size
|
|
*
|
|
* @type {?number}
|
|
*/
|
|
numContext = 1024;
|
|
|
|
/**
|
|
* Number of tokens processed in a single forward pass.
|
|
*
|
|
* @type {?number}
|
|
*/
|
|
numBatch = 1024;
|
|
|
|
/**
|
|
* Token batch size
|
|
*
|
|
* @type {?number}
|
|
*/
|
|
numUbatch = 1024;
|
|
|
|
/**
|
|
* Whether to use flash attention
|
|
*
|
|
* @type {?boolean}
|
|
*/
|
|
flashAttn = false;
|
|
|
|
/**
|
|
* Whether to use memory mapped for the model
|
|
*
|
|
* @type {?boolean}
|
|
*/
|
|
useMmap = false;
|
|
|
|
/**
|
|
* Whether to lock in memory the full model.
|
|
*
|
|
* @type {?boolean}
|
|
*/
|
|
useMlock = true;
|
|
|
|
/**
|
|
* Number of threads used during decoding.
|
|
*
|
|
* @type {?number}
|
|
*/
|
|
numThreadsDecoding = null;
|
|
|
|
/**
|
|
* The name of model file
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
modelFile = null;
|
|
|
|
/**
|
|
* The backend to use.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
backend = null;
|
|
|
|
/**
|
|
* The base URL to use for openai API requests.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
baseURL = null;
|
|
|
|
/**
|
|
* The API key to use for openai API requests.
|
|
*
|
|
* @type {?string}
|
|
*/
|
|
apiKey = null;
|
|
|
|
/**
|
|
* The options for the engine when using static embeddings.
|
|
*
|
|
* @type {?StaticEmbeddingsOptions}
|
|
*/
|
|
staticEmbeddingsOptions = null;
|
|
|
|
/**
|
|
* The service type for an OpenAIPipeline.
|
|
*
|
|
* @type {SERVICE_TYPES[keyof SERVICE_TYPES] | null}
|
|
*/
|
|
serviceType = null;
|
|
|
|
/**
|
|
* The purpose of the request, used for telemetry tracking.
|
|
*
|
|
* @type {PURPOSES[keyof PURPOSES] | null}
|
|
*/
|
|
purpose = null;
|
|
|
|
/**
|
|
* This option allows for extra headers to be passed to
|
|
* OpenAI-API-compatable endpoints
|
|
*
|
|
* @type {?Record<string, string>}
|
|
*/
|
|
extraHeaders = null;
|
|
|
|
/**
|
|
* Create a PipelineOptions instance.
|
|
*
|
|
* @param {object} options - The options for the pipeline. Must include mandatory fields.
|
|
*/
|
|
constructor(options) {
|
|
this.updateOptions(options);
|
|
}
|
|
|
|
/**
|
|
* Determines if the pipeline is mocked.
|
|
*
|
|
* It is made static to enable easier global overriding during unit tests and to allow the
|
|
* check to be performed without requiring an instance of the class.
|
|
*
|
|
* @param {object} options - The options for the pipeline.
|
|
*/
|
|
static isMocked(options) {
|
|
return (
|
|
PIPELINE_TEST_NAMES.includes(options.taskName) ||
|
|
PIPELINE_TEST_NAMES.includes(options.modelId)
|
|
);
|
|
}
|
|
|
|
/**
|
|
* Private method to validate enum fields.
|
|
*
|
|
* @param {string} field - The field being validated (e.g., 'dtype', 'device', 'executionPriority').
|
|
* @param {*} value - The value being checked against the enum.
|
|
* @throws {Error} Throws an error if the value is not valid.
|
|
* @private
|
|
*/
|
|
#validateEnum(field, value) {
|
|
const enums = {
|
|
dtype: QuantizationLevel,
|
|
device: InferenceDevice,
|
|
executionPriority: ExecutionPriority,
|
|
logLevel: LogLevel,
|
|
modelHub: ModelHub,
|
|
kvCacheDtype: KVCacheQuantizationLevel,
|
|
flashAttn: AllowedBoolean,
|
|
useMmap: AllowedBoolean,
|
|
useMlock: AllowedBoolean,
|
|
};
|
|
// Check if the value is part of the enum or null
|
|
if (!Object.values(enums[field]).includes(value)) {
|
|
throw new PipelineOptionsValidationError(field, value);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Validates the taskName field, ensuring it contains only alphanumeric characters, underscores, and dashes.
|
|
* Slashes are not allowed in the taskName.
|
|
*
|
|
* @param {string} field - The name of the field being validated (e.g., taskName).
|
|
* @param {string} value - The value of the field to validate.
|
|
* @throws {Error} Throws an error if the taskName contains invalid characters.
|
|
* @private
|
|
*/
|
|
#validateTaskName(field, value) {
|
|
// Define a regular expression to verify taskName pattern (alphanumeric, underscores, and dashes, no slashes)
|
|
const validTaskNamePattern = /^[a-zA-Z0-9_\-]+$/;
|
|
|
|
// Check if the value matches the pattern
|
|
if (!validTaskNamePattern.test(value)) {
|
|
throw new PipelineOptionsValidationError(
|
|
field,
|
|
value,
|
|
"Should contain only alphanumeric characters, underscores, or dashes."
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Validates a taskName or ID.
|
|
*
|
|
* The ID can optionally be in the form `organization/name`, where both `organization` and `name`
|
|
* follow the `taskName` pattern (alphanumeric characters, underscores, and dashes).
|
|
*
|
|
* Throws an exception if the name or ID is invalid.
|
|
*
|
|
* @param {string} field - The name of the field being validated.
|
|
* @param {string} value - The value of the field to validate.
|
|
* @throws {PipelineOptionsValidationError} Throws a validation error if the ID is invalid.
|
|
* @private
|
|
*/
|
|
#validateId(field, value) {
|
|
// Define a regular expression to match the optional organization and required name
|
|
// `organization/` part is optional, and both parts should follow the taskName pattern.
|
|
const validPattern = /^(?:[a-zA-Z0-9_\-\.\:]+\/)?[a-zA-Z0-9_\-\.\:]+$/;
|
|
|
|
// Check if the value matches the pattern
|
|
if (!validPattern.test(value)) {
|
|
throw new PipelineOptionsValidationError(
|
|
field,
|
|
value,
|
|
"Should follow the format 'organization/name' or 'name', where both parts contain only alphanumeric characters, underscores, dots or dashes."
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Generic method to validate an integer within a specified range.
|
|
*
|
|
* @param {string} field - The name of the field being validated.
|
|
* @param {number} value - The integer value to validate.
|
|
* @param {number} min - The minimum allowed value (inclusive).
|
|
* @param {number} max - The maximum allowed value (inclusive).
|
|
* @throws {Error} Throws an error if the value is not a valid integer within the range.
|
|
* @private
|
|
*/
|
|
#validateIntegerRange(field, value, min, max) {
|
|
if (!Number.isInteger(value) || value < min || value > max) {
|
|
throw new PipelineOptionsValidationError(
|
|
field,
|
|
value,
|
|
`Should be an integer between ${min} and ${max}.`
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Validates the revision field.
|
|
* The revision can be `main` or a version following a pattern like `v1.0.0`, `1.0.0-beta1`, `1.0.0.alpha2`, `1.0.0.rc1`, etc.
|
|
*
|
|
* @param {string} field - The name of the field being validated (e.g., modelRevision, tokenizerRevision).
|
|
* @param {string} value - The value of the revision field to validate.
|
|
* @throws {Error} Throws an error if the revision does not follow the expected pattern.
|
|
* @private
|
|
*/
|
|
#validateRevision(field, value) {
|
|
// Regular expression to match `main` or a version like `v1`, `v1.0.0`, `1.0.0-alpha1`, `1.0.0.alpha2`, `1.0.0.rc1`, etc.
|
|
const revisionPattern =
|
|
/^v?(\d+(\.\d+){0,2})([-\.](alpha\d*|beta\d*|pre\d*|post\d*|rc\d*))?$|^main$/;
|
|
|
|
// Check if the value matches the pattern
|
|
if (!revisionPattern.test(value)) {
|
|
throw new PipelineOptionsValidationError(
|
|
field,
|
|
value,
|
|
`Should be 'main' or follow a versioning pattern like 'v1.0.0', '1.0.0-beta1', '1.0.0.alpha2', '1.0.0.rc1', etc.`
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Validates the model file.
|
|
* The model file can be of the form filename.extension
|
|
*
|
|
* @param {string} field - The name of the field being validated (e.g., modelRevision, tokenizerRevision).
|
|
* @param {string} value - The value of the revision field to validate.
|
|
* @throws {Error} Throws an error if the model file does not follow the expected pattern.
|
|
* @private
|
|
*/
|
|
#validateModelFile(field, value) {
|
|
// Check if the value matches the pattern
|
|
if (!FILE_REGEX.test(value)) {
|
|
throw new PipelineOptionsValidationError(
|
|
field,
|
|
value,
|
|
`Should be of the form filename.extension`
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Updates multiple options at once.
|
|
*
|
|
* @param {object} options - An object containing the options to update.
|
|
* @throws {Error} Throws an error if an invalid option is provided.
|
|
*/
|
|
updateOptions(options) {
|
|
const allowedKeys = [
|
|
"engineId",
|
|
"featureId",
|
|
"taskName",
|
|
"modelHub",
|
|
"modelHubRootUrl",
|
|
"modelHubUrlTemplate",
|
|
"timeoutMS",
|
|
"modelId",
|
|
"modelRevision",
|
|
"flowId",
|
|
"tokenizerId",
|
|
"tokenizerRevision",
|
|
"processorId",
|
|
"processorRevision",
|
|
"logLevel",
|
|
"runtimeFilename",
|
|
"device",
|
|
"dtype",
|
|
"numThreads",
|
|
"executionPriority",
|
|
"useExternalDataFormat",
|
|
"kvCacheDtype",
|
|
"numContext",
|
|
"numBatch",
|
|
"numUbatch",
|
|
"flashAttn",
|
|
"useMmap",
|
|
"useMlock",
|
|
"numThreadsDecoding",
|
|
"modelFile",
|
|
"backend",
|
|
"baseURL",
|
|
"apiKey",
|
|
"staticEmbeddingsOptions",
|
|
"serviceType",
|
|
"purpose",
|
|
"extraHeaders",
|
|
];
|
|
|
|
if (options instanceof PipelineOptions) {
|
|
options = options.getOptions();
|
|
}
|
|
|
|
let optionsKeys = Object.keys(options);
|
|
|
|
allowedKeys.forEach(key => {
|
|
// If options does not have the key we can ignore it.
|
|
// We also ignore `null` values.
|
|
if (!optionsKeys.includes(key) || options[key] == null) {
|
|
return;
|
|
}
|
|
// Validating featureId
|
|
if (key === "featureId") {
|
|
const featureId = options[key];
|
|
if (FEATURES.hasOwnProperty(featureId)) {
|
|
// if featureId is set and engineId is not set, we set it
|
|
if (
|
|
options.engineId == null ||
|
|
options.engineId === DEFAULT_ENGINE_ID
|
|
) {
|
|
options.engineId = FEATURES[featureId].engineId;
|
|
this.engineId = options.engineId;
|
|
}
|
|
} else if (
|
|
// Allow tests to define a feature id.
|
|
featureId != "test-feature"
|
|
) {
|
|
// we want an explicit list of features.
|
|
throw new PipelineOptionsValidationError(
|
|
key,
|
|
featureId,
|
|
`Should be one of ${Object.keys(FEATURES).join(", ")}`
|
|
);
|
|
}
|
|
}
|
|
// Validating values.
|
|
if (key === "taskName") {
|
|
this.#validateTaskName(key, options[key]);
|
|
}
|
|
|
|
if (["modelId", "tokenizerId", "processorId"].includes(key)) {
|
|
this.#validateId(key, options[key]);
|
|
}
|
|
|
|
if (
|
|
["modelRevision", "tokenizerRevision", "processorRevision"].includes(
|
|
key
|
|
)
|
|
) {
|
|
this.#validateRevision(key, options[key]);
|
|
}
|
|
|
|
if (
|
|
[
|
|
"modelHub",
|
|
"dtype",
|
|
"device",
|
|
"executionPriority",
|
|
"logLevel",
|
|
"kvCacheDtype",
|
|
].includes(key)
|
|
) {
|
|
this.#validateEnum(key, options[key]);
|
|
}
|
|
|
|
if (["numThreads", "numThreadsDecoding"].includes(key)) {
|
|
this.#validateIntegerRange(key, options[key], 0, 100);
|
|
}
|
|
|
|
if (key === "timeoutMS") {
|
|
this.#validateIntegerRange(key, options[key], -1, 36000000);
|
|
}
|
|
|
|
if (key === "modelHub") {
|
|
ModelHub.apply(this, options[key]);
|
|
}
|
|
|
|
if (key === "modelFile") {
|
|
this.#validateModelFile(key, options[key]);
|
|
}
|
|
|
|
if (["numContext", "numBatch", "numUbatch"].includes(key)) {
|
|
this.#validateIntegerRange(key, options[key], 1, 10000000);
|
|
}
|
|
|
|
if (
|
|
key === "backend" &&
|
|
!Object.values(BACKENDS).includes(options[key]) &&
|
|
// Allow tests to define a test backend.
|
|
options[key] != "test-backend"
|
|
) {
|
|
throw new PipelineOptionsValidationError(
|
|
key,
|
|
options[key],
|
|
`Should be one of ${Object.values(BACKENDS).join(", ")}`
|
|
);
|
|
}
|
|
|
|
this[key] = options[key];
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Returns an object containing all current options.
|
|
*
|
|
* @returns {object} An object with the current options.
|
|
*/
|
|
getOptions() {
|
|
return {
|
|
engineId: this.engineId,
|
|
featureId: this.featureId,
|
|
taskName: this.taskName,
|
|
modelHub: this.modelHub,
|
|
modelHubRootUrl: this.modelHubRootUrl,
|
|
modelHubUrlTemplate: this.modelHubUrlTemplate,
|
|
timeoutMS: this.timeoutMS,
|
|
modelId: this.modelId,
|
|
modelRevision: this.modelRevision,
|
|
flowId: this.flowId,
|
|
tokenizerId: this.tokenizerId,
|
|
tokenizerRevision: this.tokenizerRevision,
|
|
processorId: this.processorId,
|
|
processorRevision: this.processorRevision,
|
|
logLevel: this.logLevel,
|
|
runtimeFilename: this.runtimeFilename,
|
|
device: this.device,
|
|
dtype: this.dtype,
|
|
numThreads: this.numThreads,
|
|
executionPriority: this.executionPriority,
|
|
useExternalDataFormat: this.useExternalDataFormat,
|
|
kvCacheDtype: this.kvCacheDtype,
|
|
numContext: this.numContext,
|
|
numBatch: this.numBatch,
|
|
numUbatch: this.numUbatch,
|
|
flashAttn: this.flashAttn,
|
|
useMmap: this.useMmap,
|
|
useMlock: this.useMlock,
|
|
numThreadsDecoding: this.numThreadsDecoding,
|
|
modelFile: this.modelFile,
|
|
backend: this.backend,
|
|
baseURL: this.baseURL,
|
|
apiKey: this.apiKey,
|
|
staticEmbeddingsOptions: this.staticEmbeddingsOptions,
|
|
serviceType: this.serviceType,
|
|
purpose: this.purpose,
|
|
extraHeaders: this.extraHeaders,
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Updates the given configuration object with the options.
|
|
*
|
|
* @param {object} config - The configuration object to be updated.
|
|
*/
|
|
applyToConfig(config) {
|
|
const options = this.getOptions();
|
|
Object.keys(options).forEach(key => {
|
|
if (options[key] !== null) {
|
|
config[key] = options[key];
|
|
}
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Per-request metadata fields that must not influence engine reuse.
|
|
* Callers differing only in these values should share one engine.
|
|
*/
|
|
static #nonIdentityKeys = new Set([
|
|
"engineId",
|
|
"featureId",
|
|
"flowId",
|
|
"logLevel",
|
|
"timeoutMS",
|
|
"serviceType",
|
|
"purpose",
|
|
]);
|
|
|
|
/**
|
|
* Checks if this PipelineOptions is equivalent to another for engine-reuse
|
|
* purposes. Fields in #nonIdentityKeys are intentionally ignored.
|
|
*
|
|
* @param {PipelineOptions} other - The other PipelineOptions instance to compare with.
|
|
* @returns {boolean} True if the instances are equal, false otherwise.
|
|
*/
|
|
equals(other) {
|
|
if (!(other instanceof PipelineOptions)) {
|
|
return false;
|
|
}
|
|
const options = this.getOptions();
|
|
const otherOptions = other.getOptions();
|
|
const skip = PipelineOptions.#nonIdentityKeys;
|
|
|
|
const isEqual = (val1, val2) => {
|
|
if (val1 === val2) {
|
|
return true;
|
|
}
|
|
if (val1 == null || val2 == null) {
|
|
return false;
|
|
}
|
|
if (typeof val1 !== "object" || typeof val2 !== "object") {
|
|
return false;
|
|
}
|
|
const keys1 = Object.keys(val1);
|
|
const keys2 = Object.keys(val2);
|
|
if (keys1.length !== keys2.length) {
|
|
return false;
|
|
}
|
|
return keys1.every(key => isEqual(val1[key], val2[key]));
|
|
};
|
|
|
|
return Object.keys(options).every(key => {
|
|
if (skip.has(key)) {
|
|
return true;
|
|
}
|
|
if (key === "backend") {
|
|
return PipelineOptions.#backendsCompatible(
|
|
options.backend,
|
|
otherOptions.backend
|
|
);
|
|
}
|
|
return isEqual(options[key], otherOptions[key]);
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Whether two backend identifiers should be considered equivalent for
|
|
* engine reuse. "best-onnx" is a sentinel that resolves at engine-creation
|
|
* time to either "onnx" or "onnx-native"; a cached engine using either
|
|
* concrete backend should still satisfy a new "best-onnx" request.
|
|
*/
|
|
static #backendsCompatible(a, b) {
|
|
if (a === b) {
|
|
return true;
|
|
}
|
|
const isOnnxConcrete = v => v === "onnx" || v === "onnx-native";
|
|
if (a === "best-onnx" && isOnnxConcrete(b)) {
|
|
return true;
|
|
}
|
|
if (b === "best-onnx" && isOnnxConcrete(a)) {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* This class controls the life cycle of the engine process used both in the
|
|
* Translations engine and the MLEngine component.
|
|
*/
|
|
export class EngineProcess {
|
|
/**
|
|
* Get a reference to all running "inference" processes.
|
|
*
|
|
* @returns {sequence<nsIDOMProcessParent>}
|
|
*/
|
|
static #inferenceProcesses() {
|
|
return ChromeUtils.getAllDOMProcesses().filter(
|
|
p => p.remoteType == "inference"
|
|
);
|
|
}
|
|
|
|
/**
|
|
* See if all engines are terminated and the "inference" process has been shut
|
|
* down. This is useful for testing.
|
|
*
|
|
* @returns {boolean}
|
|
*/
|
|
static areAllEnginesTerminated() {
|
|
return !EngineProcess.#inferenceProcesses().length;
|
|
}
|
|
|
|
/**
|
|
* @returns {Promise<TranslationsEngineParent>}
|
|
*/
|
|
static async getTranslationsEngineParent() {
|
|
return EngineProcess.#getEngineActor({ actorName: "TranslationsEngine" });
|
|
}
|
|
|
|
/**
|
|
* @returns {Promise<MLEngineParent>}
|
|
*/
|
|
static async getMLEngineParent() {
|
|
// the pref is off by default
|
|
if (!Services.prefs.getBoolPref("browser.ml.enable")) {
|
|
throw new Error("MLEngine is disabled. Check the browser.ml prefs.");
|
|
}
|
|
|
|
return EngineProcess.#getEngineActor({ actorName: "MLEngine" });
|
|
}
|
|
|
|
/**
|
|
* @returns {Promise<JSProcessActorParent>}
|
|
*/
|
|
static async #getEngineActor({ actorName }) {
|
|
let keepAlive = await ChromeUtils.ensureHeadlessContentProcess(
|
|
"inference",
|
|
{ preferUsed: true }
|
|
);
|
|
if (!keepAlive?.domProcess?.canSend) {
|
|
return null;
|
|
}
|
|
|
|
try {
|
|
const actor = keepAlive.domProcess.getActor(actorName);
|
|
|
|
// keep track of the childID for the inference process, so we can observe its shutdowns.
|
|
actor.childID = keepAlive.domProcess.childID;
|
|
|
|
if (actor && !actor.processKeepAlive) {
|
|
ChromeUtils.addProfilerMarker(
|
|
"EngineProcess",
|
|
{},
|
|
`Setting ${actorName} "inference" process keep-alive`
|
|
);
|
|
actor.processKeepAlive = keepAlive;
|
|
keepAlive = null;
|
|
}
|
|
return actor;
|
|
} finally {
|
|
if (keepAlive) {
|
|
keepAlive.invalidateKeepAlive();
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Send the `ForceShutdown` message to the TranslationsEngine, terminating
|
|
* running engines, and potentially leading to "inference" process shutdown.
|
|
*/
|
|
static destroyTranslationsEngine() {
|
|
return EngineProcess.#forceShutdownEngine({
|
|
actorName: "TranslationsEngine",
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Send the `ForceShutdown` message to the MLEngine, terminating running
|
|
* queries, and potentially leading to "inference" process shutdown.
|
|
*/
|
|
static destroyMLEngine() {
|
|
return EngineProcess.#forceShutdownEngine({ actorName: "MLEngine" });
|
|
}
|
|
|
|
static #forceShutdownEngine({ actorName }) {
|
|
return Promise.allSettled(
|
|
EngineProcess.#inferenceProcesses().map(async process => {
|
|
let actor = process.getExistingActor(actorName);
|
|
if (actor) {
|
|
await actor.forceShutdown();
|
|
|
|
// The actor should have cleared its own KeepAlive.
|
|
if (actor.processKeepAlive) {
|
|
ChromeUtils.addProfilerMarker(
|
|
"EngineProcess",
|
|
{},
|
|
`Force-dropping ${actorName} "inference" process keep-alive`
|
|
);
|
|
actor.processKeepAlive.invalidateKeepAlive();
|
|
actor.processKeepAlive = null;
|
|
}
|
|
}
|
|
})
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Creates a new `MLEngine` instance with the provided options.
|
|
*
|
|
* @template {EngineFeatureIds} FeatureID
|
|
*
|
|
* @param {object} options - Configuration options for the ML engine.
|
|
* @param {?function(ProgressAndStatusCallbackParams):void} [notificationsCallback] - A function to call to indicate notifications.
|
|
* @param {?AbortSignal} [abortSignal] - AbortSignal to cancel the download.
|
|
* @returns {Promise<MLEngine<FeatureID>>}
|
|
*/
|
|
export async function createEngine(
|
|
options,
|
|
notificationsCallback = null,
|
|
abortSignal
|
|
) {
|
|
try {
|
|
const pipelineOptions = new PipelineOptions(options);
|
|
const engineParent = await EngineProcess.getMLEngineParent();
|
|
return engineParent.getEngine({
|
|
pipelineOptions,
|
|
notificationsCallback,
|
|
abortSignal,
|
|
});
|
|
} catch (e) {
|
|
Glean.firefoxAiRuntime.engineCreationFailure.record({
|
|
engineId: options.engineId || "",
|
|
modelId: options.modelId || "",
|
|
featureId: options.featureId || "",
|
|
flow_id: options.flowId || "",
|
|
taskName: options.taskName || "",
|
|
error: e.constructor.name || "",
|
|
});
|
|
throw e;
|
|
}
|
|
}
|