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
// TextGenerator is a ChromeOnly WebIDL interface (dom/chrome-webidl).
/* global TextGenerator */
/** MLEngine-shaped engine over the HWInference utility process. */
import { XPCOMUtils } from "resource://gre/modules/XPCOMUtils.sys.mjs";
/**
* @typedef {import("chrome://global/content/ml/EngineProcess.sys.mjs").PipelineOptions} PipelineOptions
* @typedef {import("chrome://global/content/ml/Utils.sys.mjs").ProgressAndStatusCallbackParams} ProgressAndStatusCallbackParams
* @typedef {(data: ProgressAndStatusCallbackParams) => void} NotificationsCallback
* @typedef {{ name: string, when: number }} RunTimestamp
*/
/**
* PipelineOptions declares its model fields nullable; createEngine() only
* hands this engine options it has resolved.
*
* @typedef {PipelineOptions & {
* engineId: string,
* taskName: string,
* modelId: string,
* modelRevision: string,
* modelFile: string,
* modelHubRootUrl: string,
* modelHubUrlTemplate: string,
* }} ResolvedPipelineOptions
*/
/**
* The llama.cpp request shape MLEngine consumers send to run().
*
* @typedef {object} LlamaRunRequest
* @property {string | TextGenerationMessage[]} prompt
* @property {number} [nPredict]
* @property {number} [minOutputBufferSize]
* @property {TextGenerationSampler[]} [samplers]
* @property {number[]} [stopTokens]
* @property {boolean} [stopOnEndOfGenerationTokens]
*/
/**
* @typedef {object} EngineRunResult
* @property {true} done
* @property {string} finalOutput
* @property {true} ok
* @property {ReturnType<typeof toEngineMetrics>} metrics
* @property {ReturnType<typeof toEngineResources>} resourcesBefore
* @property {ReturnType<typeof toEngineResources>} resourcesAfter
*/
const lazy = XPCOMUtils.declareLazy({
LLAMA_CPP_VERSION:
"moz-src:///toolkit/components/ml/actors/MLEngineParent.sys.mjs",
MLEngineParent:
"moz-src:///toolkit/components/ml/actors/MLEngineParent.sys.mjs",
MLTelemetry: "chrome://global/content/ml/MLTelemetry.sys.mjs",
generateUUID: "chrome://global/content/ml/Utils.sys.mjs",
});
/** @param {TextGenerationResourceSnapshot} snapshot */
const toEngineResources = snapshot => ({
cpuTime: snapshot.cpuTimeMs,
memory: snapshot.memoryBytes,
});
// LlamaCppPipeline spells the float cache types fp16 and fp32; the
// generator, f16 and f32.
/**
* @param {string} dtype
* @returns {TextGenerationKVCacheDtype}
*/
function toKVCacheDtype(dtype) {
return /** @type {TextGenerationKVCacheDtype} */ (dtype.replace(/^fp/, "f"));
}
/**
* @param {PipelineOptions} options
* @returns {TextGeneratorCreateOptions}
*/
function toCreateOptions(options) {
/** @type {TextGeneratorCreateOptions} */
const createOptions = {
contextSize: options.numContext ?? undefined,
batchSize: options.numBatch ?? undefined,
ubatchSize: options.numUbatch ?? undefined,
flashAttn: options.flashAttn ?? undefined,
featureId: options.featureId ?? "",
};
if (options.kvCacheDtype) {
createOptions.kvCacheDtype = toKVCacheDtype(options.kvCacheDtype);
}
if (options.numThreads) {
createOptions.numThreads = options.numThreads;
}
if (options.numThreadsDecoding) {
createOptions.numThreadsDecoding = options.numThreadsDecoding;
}
return createOptions;
}
/**
* @param {LlamaRunRequest} request
* @returns {TextGenerationRequest}
*/
function toGenerateRequest(request) {
let prompt = request.prompt;
if (!Array.isArray(prompt)) {
prompt = [{ role: "user", content: String(prompt) }];
}
/** @type {TextGenerationRequest} */
const generateRequest = {
messages: prompt.map(message => ({
role: message.role,
content: message.content,
})),
};
// Leave unset fields absent so TextGenerationRequest defaults apply.
if (request.nPredict != null) {
generateRequest.maxTokens = request.nPredict;
}
if (request.minOutputBufferSize != null) {
generateRequest.bufferLength = request.minOutputBufferSize;
}
if (request.samplers != null) {
generateRequest.samplers = request.samplers;
}
if (request.stopTokens != null) {
generateRequest.stopTokens = request.stopTokens;
}
if (request.stopOnEndOfGenerationTokens != null) {
generateRequest.stopOnEndOfGenerationTokens =
request.stopOnEndOfGenerationTokens;
}
return generateRequest;
}
/**
* @param {TextGenerationResult} result
* @param {RunTimestamp[]} runTimestamps
*/
function toEngineMetrics(result, runTimestamps) {
const { usage } = result;
const outputTokens = usage.generatedTokens;
const decodingTime = usage.timings.decodeMs;
return {
runTimestamps,
inputTokens: usage.promptTokens,
inputCharacters: usage.promptCharacters,
outputTokens,
inferenceTime: usage.timings.prefillMs + decodingTime,
decodingTime,
timeToFirstToken: usage.timings.prefillMs,
tokensPerSecond: decodingTime
? outputTokens / (decodingTime / 1000)
: undefined,
timePerOutputToken: outputTokens ? decodingTime / outputTokens : undefined,
};
}
/** Async-iterable sink that the delta callback pushes into. */
class ChunkQueue {
/** @type {string[]} */
#pending = [];
#closed = false;
/** @type {((value?: unknown) => void) | null} */
#wakeUp = null;
/** @param {string} text */
push(text) {
this.#pending.push(text);
this.#wakeUp?.();
}
close() {
this.#closed = true;
this.#wakeUp?.();
}
async *[Symbol.asyncIterator]() {
while (!this.#closed || this.#pending.length) {
if (!this.#pending.length) {
await new Promise(resolve => {
this.#wakeUp = resolve;
});
this.#wakeUp = null;
continue;
}
yield this.#pending.shift();
}
}
}
/** What recordEngineRun reports about a streamed run. */
class StreamStats {
firstChunkAt = 0;
lastChunkAt = 0;
chunkCount = 0;
characterCount = 0;
streaming = false;
constructor() {
this.beforeRun = ChromeUtils.now();
}
/** @param {string} text */
onChunk(text) {
this.lastChunkAt = ChromeUtils.now();
if (!this.firstChunkAt) {
this.firstChunkAt = this.lastChunkAt;
}
this.chunkCount += 1;
this.characterCount += text.length;
}
/** @param {EngineRunResult} result */
metrics(result) {
return {
tokenCount: result.metrics.outputTokens,
characterCount: this.characterCount,
timeToFirstChunk: this.firstChunkAt
? this.firstChunkAt - this.beforeRun
: null,
averageChunkTime:
this.chunkCount > 1
? (this.lastChunkAt - this.firstChunkAt) / (this.chunkCount - 1)
: null,
};
}
}
export class TextGenerationEngine {
/** @type {string} */
engineId;
/** @type {ResolvedPipelineOptions} */
pipelineOptions;
/** @type {"uninitialized" | "ready" | "closed"} */
engineStatus = "uninitialized";
/** @type {NotificationsCallback | null} */
notificationsCallback = null;
/** @type {InstanceType<typeof lazy.MLTelemetry>} */
telemetry;
/** @type {TextGenerator | null} */
#generator = null;
/** @type {RunTimestamp[]} */
#initTimestamps = [];
#inFlight = false;
/** @param {PipelineOptions} pipelineOptions */
static shouldRoute(pipelineOptions) {
if (!Services.prefs.getBoolPref("browser.ml.llama.hwInference", false)) {
return false;
}
return pipelineOptions.backend === "llama.cpp";
}
/**
* @param {PipelineOptions} pipelineOptions
* @param {NotificationsCallback | null} [notificationsCallback]
* @param {AbortSignal} [abortSignal]
*/
static async create(
pipelineOptions,
notificationsCallback = null,
abortSignal = undefined
) {
const start = ChromeUtils.now();
const engine = new TextGenerationEngine(
pipelineOptions,
notificationsCallback
);
const { engineId } = engine;
try {
await engine.#initialize(abortSignal);
engine.telemetry.recordEngineCreationSuccessFlow({
engineId,
duration: ChromeUtils.now() - start,
});
} catch (error) {
engine.telemetry.recordEngineCreationFailure({
modelId: pipelineOptions.modelId,
featureId: pipelineOptions.featureId,
taskName: pipelineOptions.taskName,
engineId,
error,
});
throw error;
}
return engine;
}
/**
* @param {PipelineOptions} pipelineOptions
* @param {NotificationsCallback | null} notificationsCallback
*/
constructor(pipelineOptions, notificationsCallback) {
this.pipelineOptions = /** @type {ResolvedPipelineOptions} */ (
pipelineOptions
);
this.engineId = this.pipelineOptions.engineId;
this.notificationsCallback = notificationsCallback;
this.telemetry = new lazy.MLTelemetry({
featureId: pipelineOptions.featureId,
flowId: pipelineOptions.flowId,
hostProcess: "hwinference",
});
}
/** @param {AbortSignal} [abortSignal] */
async #initialize(abortSignal) {
abortSignal?.throwIfAborted();
const options = this.pipelineOptions;
this.#initTimestamps = [
{ name: "initializationStart", when: ChromeUtils.now() },
];
const hub = await lazy.MLEngineParent.createModelHub({
rootUrl: options.modelHubRootUrl,
urlTemplate: options.modelHubUrlTemplate,
});
const sessionId = lazy.generateUUID();
let modelBlob;
try {
[modelBlob] = await hub.getModelFileAsBlob({
engineId: options.engineId,
taskName: options.taskName,
model: options.modelId,
revision: options.modelRevision,
file: options.modelFile,
modelHubRootUrl: options.modelHubRootUrl,
modelHubUrlTemplate: options.modelHubUrlTemplate,
progressCallback: this.notificationsCallback,
abortSignal,
featureId: options.featureId ?? options.engineId,
sessionId,
});
} finally {
await hub.notifyModelDownloadComplete({
engineId: options.engineId,
model: options.modelId,
revision: options.modelRevision,
featureId: options.featureId ?? options.engineId,
sessionId,
});
}
const createOptions = toCreateOptions(options);
this.#generator = await TextGenerator.create(modelBlob, createOptions);
this.pipelineOptions.backend = "llama.cpp";
this.#initTimestamps.push({
name: "initializationEnd",
when: ChromeUtils.now(),
});
this.engineStatus = "ready";
}
#runnableGenerator() {
if (this.engineStatus !== "ready" || !this.#generator) {
throw new Error("TextGenerationEngine: engine is not ready");
}
if (this.#inFlight) {
throw new Error("A generation is already in progress");
}
return this.#generator;
}
/**
* @param {LlamaRunRequest} request
* @param {(text: string) => void} [onDelta]
* @returns {Promise<EngineRunResult>}
*/
async #execute(request, onDelta) {
const generator = this.#runnableGenerator();
this.#inFlight = true;
try {
const runStart = ChromeUtils.now();
generator.clear();
const result = await generator.generate(
toGenerateRequest(request),
onDelta
);
const runEnd = ChromeUtils.now();
return {
done: true,
finalOutput: result.content,
ok: true,
metrics: toEngineMetrics(result, [
...this.#initTimestamps,
{ name: "runStart", when: runStart },
{ name: "runEnd", when: runEnd },
]),
resourcesBefore: toEngineResources(result.resources.before),
resourcesAfter: toEngineResources(result.resources.after),
};
} finally {
this.#inFlight = false;
}
}
/**
* @param {StreamStats} stats
* @param {EngineRunResult} result
*/
#recordRun(stats, result) {
this.telemetry.recordRunInferenceSuccessFlow(this.engineId, result.metrics);
this.telemetry.recordEngineRun({
beforeRun: stats.beforeRun,
resourcesBefore: result.resourcesBefore,
resourcesAfter: result.resourcesAfter,
engineId: this.engineId,
modelId: this.pipelineOptions.modelId,
backend: this.pipelineOptions.backend,
backendSourceRevision: lazy.LLAMA_CPP_VERSION,
...(stats.streaming ? stats.metrics(result) : {}),
});
}
/** @param {LlamaRunRequest} request */
async run(request) {
const stats = new StreamStats();
try {
const result = await this.#execute(request);
this.#recordRun(stats, result);
return result;
} catch (error) {
this.telemetry.recordRunInferenceFailure(error);
throw error;
}
}
/** @param {LlamaRunRequest} request */
async *runWithGenerator(request) {
const stats = new StreamStats();
stats.streaming = true;
const queue = new ChunkQueue();
/** @type {Promise<EngineRunResult> | null} */
let completion = null;
let settled = false;
try {
completion = this.#execute(request, text => {
stats.onChunk(text);
queue.push(text);
});
// Keeps the rejection handled; the await below rethrows it.
completion.catch(() => {});
const close = () => queue.close();
completion.then(close, close);
for await (const text of queue) {
yield { text, tokens: [], isPrompt: false };
}
const result = await completion;
// Link Preview keys its final flush on an empty terminal chunk.
yield { text: "", tokens: [], isPrompt: false };
this.#recordRun(stats, result);
settled = true;
return result;
} catch (error) {
settled = true;
this.telemetry.recordRunInferenceFailure(error);
throw error;
} finally {
// Breaking out of a `for await` loop returns this generator mid-decode.
if (!settled) {
this.cancel();
const result = await completion?.catch(() => null);
if (result) {
this.#recordRun(stats, result);
}
}
}
}
/** Stops the generation; the run settles and the generator stays usable. */
cancel() {
this.#generator?.cancel();
}
async terminate() {
if (this.#generator) {
this.#generator.terminate();
this.#generator = null;
}
this.engineStatus = "closed";
}
}