chore: import upstream snapshot with attribution
Build / build (push) Has been cancelled
Tests / test (push) Has been cancelled
Build site and push to gh-pages / Build site (push) Has been cancelled
Linter / lint (push) Has been cancelled
Security / dependency-review (push) Has been cancelled
Security / npm-audit (push) Has been cancelled
Security / codeql (push) Has been cancelled
Build / build (push) Has been cancelled
Tests / test (push) Has been cancelled
Build site and push to gh-pages / Build site (push) Has been cancelled
Linter / lint (push) Has been cancelled
Security / dependency-review (push) Has been cancelled
Security / npm-audit (push) Has been cancelled
Security / codeql (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,194 @@
|
||||
import * as tvmjs from "@mlc-ai/web-runtime";
|
||||
import {
|
||||
AppConfig,
|
||||
ChatConfig,
|
||||
ModelRecord,
|
||||
prebuiltAppConfig,
|
||||
getCacheBackend,
|
||||
} from "./config";
|
||||
import { cleanModelUrl } from "./support";
|
||||
import { ModelNotFoundError, UnsupportedTokenizerFilesError } from "./error";
|
||||
import { Tokenizer } from "@mlc-ai/web-tokenizers";
|
||||
import { ModelIntegrity, verifyIntegrity } from "./integrity";
|
||||
|
||||
type CacheScope = "webllm/model" | "webllm/config" | "webllm/wasm";
|
||||
type CacheOptions = Pick<
|
||||
tvmjs.TensorCacheAccessOptions,
|
||||
"cacheType" | "opfsAccessMode"
|
||||
>;
|
||||
|
||||
export function getCacheOptions(appConfig: AppConfig): CacheOptions {
|
||||
const options: CacheOptions = {
|
||||
cacheType: getCacheBackend(appConfig),
|
||||
};
|
||||
if (appConfig.opfsAccessMode !== undefined) {
|
||||
options.opfsAccessMode = appConfig.opfsAccessMode;
|
||||
}
|
||||
return options;
|
||||
}
|
||||
|
||||
export function getTensorCacheAccessOptions(
|
||||
scope: CacheScope,
|
||||
appConfig: AppConfig,
|
||||
): tvmjs.TensorCacheAccessOptions {
|
||||
return {
|
||||
cacheScope: scope,
|
||||
...getCacheOptions(appConfig),
|
||||
};
|
||||
}
|
||||
|
||||
function createScopedArtifactCache(
|
||||
scope: CacheScope,
|
||||
appConfig: AppConfig,
|
||||
): tvmjs.ArtifactCacheTemplate {
|
||||
return tvmjs.createArtifactCache(scope, getCacheOptions(appConfig));
|
||||
}
|
||||
|
||||
async function maybeVerifyTokenizerIntegrity(
|
||||
data: ArrayBuffer,
|
||||
filename: string,
|
||||
url: string,
|
||||
integrity?: ModelIntegrity,
|
||||
): Promise<void> {
|
||||
const hash = integrity?.tokenizer?.[filename];
|
||||
if (hash) {
|
||||
await verifyIntegrity(data, hash, url, integrity?.onFailure);
|
||||
}
|
||||
}
|
||||
|
||||
function findModelRecord(modelId: string, appConfig?: AppConfig): ModelRecord {
|
||||
const matchedItem = appConfig?.model_list.find(
|
||||
(item) => item.model_id == modelId,
|
||||
);
|
||||
if (matchedItem !== undefined) {
|
||||
return matchedItem;
|
||||
}
|
||||
throw new ModelNotFoundError(modelId);
|
||||
}
|
||||
|
||||
export async function hasModelInCache(
|
||||
modelId: string,
|
||||
appConfig?: AppConfig,
|
||||
): Promise<boolean> {
|
||||
if (appConfig === undefined) {
|
||||
appConfig = prebuiltAppConfig;
|
||||
}
|
||||
const modelRecord = findModelRecord(modelId, appConfig);
|
||||
const modelUrl = cleanModelUrl(modelRecord.model);
|
||||
return tvmjs.hasTensorInCache(
|
||||
modelUrl,
|
||||
getTensorCacheAccessOptions("webllm/model", appConfig),
|
||||
);
|
||||
}
|
||||
|
||||
export async function deleteModelAllInfoInCache(
|
||||
modelId: string,
|
||||
appConfig?: AppConfig,
|
||||
) {
|
||||
// function to delete model all information in cache
|
||||
if (appConfig === undefined) {
|
||||
appConfig = prebuiltAppConfig;
|
||||
}
|
||||
// delete model and tokenizer in Cache
|
||||
await deleteModelInCache(modelId, appConfig);
|
||||
// delete wasm in cache
|
||||
await deleteModelWasmInCache(modelId, appConfig);
|
||||
// delete chat config
|
||||
await deleteChatConfigInCache(modelId, appConfig);
|
||||
}
|
||||
|
||||
export async function deleteModelInCache(
|
||||
modelId: string,
|
||||
appConfig?: AppConfig,
|
||||
) {
|
||||
// delete the model NDArray In Cache
|
||||
if (appConfig === undefined) {
|
||||
appConfig = prebuiltAppConfig;
|
||||
}
|
||||
const modelRecord = findModelRecord(modelId, appConfig);
|
||||
const modelUrl = cleanModelUrl(modelRecord.model);
|
||||
const modelCache = createScopedArtifactCache("webllm/model", appConfig);
|
||||
await tvmjs.deleteTensorCache(
|
||||
modelUrl,
|
||||
getTensorCacheAccessOptions("webllm/model", appConfig),
|
||||
);
|
||||
await modelCache.deleteInCache(new URL("tokenizer.model", modelUrl).href);
|
||||
await modelCache.deleteInCache(new URL("tokenizer.json", modelUrl).href);
|
||||
}
|
||||
|
||||
export async function deleteChatConfigInCache(
|
||||
modelId: string,
|
||||
appConfig?: AppConfig,
|
||||
) {
|
||||
// delete the chat configuration in Cache
|
||||
if (appConfig === undefined) {
|
||||
appConfig = prebuiltAppConfig;
|
||||
}
|
||||
const modelRecord = findModelRecord(modelId, appConfig);
|
||||
const configCache = createScopedArtifactCache("webllm/config", appConfig);
|
||||
const modelUrl = cleanModelUrl(modelRecord.model);
|
||||
const configUrl = new URL("mlc-chat-config.json", modelUrl).href;
|
||||
await configCache.deleteInCache(configUrl);
|
||||
}
|
||||
|
||||
export async function deleteModelWasmInCache(
|
||||
modelId: string,
|
||||
appConfig?: AppConfig,
|
||||
) {
|
||||
// delete the wasm in Cache
|
||||
if (appConfig === undefined) {
|
||||
appConfig = prebuiltAppConfig;
|
||||
}
|
||||
const modelRecord = findModelRecord(modelId, appConfig);
|
||||
const wasmCache = createScopedArtifactCache("webllm/wasm", appConfig);
|
||||
await wasmCache.deleteInCache(modelRecord.model_lib);
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param baseUrl The link to which we can find tokenizer files, usually is a `ModelRecord.model`.
|
||||
* @param config A ChatConfig, usually loaded from `mlc-chat-config.json` in `baseUrl`.
|
||||
* @param appConfig An AppConfig, usually `webllm.prebuiltAppConfig` if not defined by user.
|
||||
* @param logger Logging function, console.log by default.
|
||||
* @param integrity Optional integrity configuration for verifying tokenizer files.
|
||||
* @returns
|
||||
*/
|
||||
export async function asyncLoadTokenizer(
|
||||
baseUrl: string,
|
||||
config: ChatConfig,
|
||||
appConfig: AppConfig,
|
||||
logger: (msg: string) => void = console.log,
|
||||
integrity?: ModelIntegrity,
|
||||
): Promise<Tokenizer> {
|
||||
const modelCache = createScopedArtifactCache("webllm/model", appConfig);
|
||||
|
||||
if (config.tokenizer_files.includes("tokenizer.json")) {
|
||||
const url = new URL("tokenizer.json", baseUrl).href;
|
||||
const model = await modelCache.fetchWithCache(url, "arraybuffer");
|
||||
await maybeVerifyTokenizerIntegrity(
|
||||
model,
|
||||
"tokenizer.json",
|
||||
url,
|
||||
integrity,
|
||||
);
|
||||
return Tokenizer.fromJSON(model);
|
||||
} else if (config.tokenizer_files.includes("tokenizer.model")) {
|
||||
logger(
|
||||
"Using `tokenizer.model` since we cannot locate `tokenizer.json`.\n" +
|
||||
"It is recommended to use `tokenizer.json` to ensure all token mappings are included, " +
|
||||
"since currently, files like `added_tokens.json`, `tokenizer_config.json` are ignored.\n" +
|
||||
"Consider converting `tokenizer.model` to `tokenizer.json` by compiling the model " +
|
||||
"with MLC again, or see if MLC's huggingface provides this file.",
|
||||
);
|
||||
const url = new URL("tokenizer.model", baseUrl).href;
|
||||
const model = await modelCache.fetchWithCache(url, "arraybuffer");
|
||||
await maybeVerifyTokenizerIntegrity(
|
||||
model,
|
||||
"tokenizer.model",
|
||||
url,
|
||||
integrity,
|
||||
);
|
||||
return Tokenizer.fromSentencePiece(model);
|
||||
}
|
||||
throw new UnsupportedTokenizerFilesError(config.tokenizer_files);
|
||||
}
|
||||
+2603
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,568 @@
|
||||
import {
|
||||
ChatConfig,
|
||||
ConvTemplateConfig,
|
||||
MessagePlaceholders,
|
||||
Role,
|
||||
} from "./config";
|
||||
import {
|
||||
ChatCompletionContentPart,
|
||||
ChatCompletionContentPartImage,
|
||||
ChatCompletionMessageParam,
|
||||
ChatCompletionRequest,
|
||||
} from "./openai_api_protocols/index";
|
||||
import {
|
||||
ContentTypeError,
|
||||
FunctionNotFoundError,
|
||||
InvalidToolChoiceError,
|
||||
MessageOrderError,
|
||||
MultipleTextContentError,
|
||||
SystemMessageOrderError,
|
||||
TextCompletionConversationError,
|
||||
TextCompletionConversationExpectsPrompt,
|
||||
UnsupportedRoleError,
|
||||
UnsupportedToolChoiceTypeError,
|
||||
UnsupportedToolTypeError,
|
||||
} from "./error";
|
||||
|
||||
type ImageURL = ChatCompletionContentPartImage.ImageURL;
|
||||
|
||||
/**
|
||||
* Helper to keep track of history conversations.
|
||||
*/
|
||||
export class Conversation {
|
||||
// NOTE: Update `compareConversationObject()` whenever a new state is introduced.
|
||||
/** Each message is a tuple of (Role, role_name_str, message), where message can be either a
|
||||
* string or an array of contentPart for possible image input.
|
||||
*/
|
||||
public messages: Array<
|
||||
[Role, string, string | Array<ChatCompletionContentPart> | undefined]
|
||||
> = [];
|
||||
readonly config: ConvTemplateConfig;
|
||||
|
||||
/** Whether the Conversation object is for text completion with no conversation-style formatting */
|
||||
public isTextCompletion: boolean;
|
||||
/** Used when isTextCompletion is true */
|
||||
public prompt: string | undefined;
|
||||
|
||||
public function_string = "";
|
||||
public use_function_calling = false;
|
||||
public override_system_message?: string = undefined;
|
||||
|
||||
/**
|
||||
* Tracks whether the last message is an empty thinking block. Should only
|
||||
* be true when we are in the middle of a generation. Will be set to
|
||||
* false when the reply is finished with `finishReply()`.
|
||||
*/
|
||||
private isLastMessageEmptyThinkingReplyHeader = false;
|
||||
|
||||
// TODO(tvm-team) confirm and remove
|
||||
// private contextWindowStart = 0;
|
||||
|
||||
constructor(config: ConvTemplateConfig, isTextCompletion = false) {
|
||||
this.config = config;
|
||||
this.isTextCompletion = isTextCompletion;
|
||||
}
|
||||
|
||||
// TODO: Consider rewriting this method, a bit messy.
|
||||
private getPromptArrayInternal(
|
||||
addSystem: boolean,
|
||||
startPos: number,
|
||||
config?: ChatConfig,
|
||||
): Array<string | Array<string | ImageURL>> {
|
||||
if (this.config.seps.length == 0) {
|
||||
throw Error("Need seps to work");
|
||||
}
|
||||
|
||||
// Prepare system message
|
||||
// Get overridden system message if exists, else use default one in config
|
||||
let system_message = this.config.system_message;
|
||||
if (this.override_system_message !== undefined) {
|
||||
system_message = this.override_system_message;
|
||||
}
|
||||
const system_prompt = this.config.system_template.replace(
|
||||
MessagePlaceholders.system,
|
||||
system_message,
|
||||
);
|
||||
const ret: Array<string | Array<string | ImageURL>> =
|
||||
addSystem && system_prompt !== "" ? [system_prompt] : [];
|
||||
|
||||
// Process each message in this.messages
|
||||
for (let i = startPos; i < this.messages.length; ++i) {
|
||||
const item = this.messages[i];
|
||||
const role = item[0];
|
||||
const role_str = item[1];
|
||||
const messageContent = item[2];
|
||||
|
||||
// 1. Message from `appendReplyHeader()`, message is empty; not much processing is needed.
|
||||
if (messageContent === undefined) {
|
||||
if (i !== this.messages.length - 1) {
|
||||
throw new Error(
|
||||
"InternalError: Only expect message to be undefined for last " +
|
||||
"message for a reply header.",
|
||||
);
|
||||
}
|
||||
// Add ": " if there is no such field. If "", do not add sep
|
||||
const empty_sep =
|
||||
this.config.role_empty_sep || this.config.role_empty_sep == ""
|
||||
? this.config.role_empty_sep
|
||||
: ": ";
|
||||
ret.push(role_str + empty_sep);
|
||||
continue;
|
||||
}
|
||||
|
||||
// 2. Message from `appendEmptyThinkingReplyHeader()`, message is an empty thinking block.
|
||||
if (
|
||||
this.isLastMessageEmptyThinkingReplyHeader &&
|
||||
i === this.messages.length - 1
|
||||
) {
|
||||
// TODO(Charlie): content_sep or empty_sep? For Qwen3, both are "\n".
|
||||
const content_sep =
|
||||
this.config.role_content_sep || this.config.role_content_sep == ""
|
||||
? this.config.role_content_sep
|
||||
: ": ";
|
||||
ret.push(role_str + content_sep + messageContent);
|
||||
continue;
|
||||
}
|
||||
|
||||
// 3. Each messageContent consists of one textPart, and >= 0 imageParts, regardless whether
|
||||
// it is Array<ChatCompletionContentPart> or text message. So we extract out each.
|
||||
let textContentPart = ""; // if no textPart, use an empty string
|
||||
const imageContentParts: ImageURL[] = [];
|
||||
if (Array.isArray(messageContent)) {
|
||||
// 2.1 content is Array<ChatCompletionContentPart>
|
||||
// Iterate through the contentParts, get the text and list of images. There should
|
||||
// be only a single text. TODO: is it always the case the number of textContentPart <= 1?
|
||||
let seenText = false;
|
||||
for (let i = 0; i < messageContent.length; i++) {
|
||||
const curContentPart = messageContent[i];
|
||||
if (curContentPart.type === "text") {
|
||||
if (seenText) {
|
||||
throw new MultipleTextContentError();
|
||||
}
|
||||
textContentPart = curContentPart.text;
|
||||
seenText = true;
|
||||
} else {
|
||||
imageContentParts.push(curContentPart.image_url);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 2.2 content is just a string
|
||||
textContentPart = messageContent;
|
||||
}
|
||||
|
||||
// 3. Format textContentPart with role and sep to get message_str and role_prefix
|
||||
let message_str;
|
||||
let role_prefix;
|
||||
if (this.config.role_templates !== undefined) {
|
||||
message_str = this.config.role_templates[role]?.replace(
|
||||
MessagePlaceholders[Role[role] as keyof typeof MessagePlaceholders],
|
||||
textContentPart,
|
||||
);
|
||||
if (this.use_function_calling && this.function_string !== "") {
|
||||
message_str = message_str?.replace(
|
||||
MessagePlaceholders.function,
|
||||
this.function_string,
|
||||
);
|
||||
}
|
||||
message_str = message_str?.replace(MessagePlaceholders.function, "");
|
||||
}
|
||||
|
||||
if (message_str == undefined) {
|
||||
message_str = textContentPart;
|
||||
}
|
||||
if (
|
||||
this.config.add_role_after_system_message === false &&
|
||||
system_prompt != "" &&
|
||||
i == 0
|
||||
) {
|
||||
role_prefix = "";
|
||||
} else {
|
||||
// Add ": " if there is no such field. If "", do not add sep
|
||||
const content_sep =
|
||||
this.config.role_content_sep || this.config.role_content_sep == ""
|
||||
? this.config.role_content_sep
|
||||
: ": ";
|
||||
role_prefix = role_str + content_sep;
|
||||
}
|
||||
|
||||
// 4. Combine everything together
|
||||
if (imageContentParts.length === 0) {
|
||||
// If no image, just a single string to represent this message
|
||||
ret.push(
|
||||
role_prefix +
|
||||
message_str +
|
||||
this.config.seps[i % this.config.seps.length],
|
||||
);
|
||||
} else {
|
||||
// Per-model image layout
|
||||
const curMessageList: Array<string | ImageURL> = [role_prefix];
|
||||
const modelType = config?.model_type ?? "";
|
||||
imageContentParts.forEach((curImage: ImageURL) => {
|
||||
if (modelType === "phi3_v") {
|
||||
curMessageList.push(curImage);
|
||||
curMessageList.push("\n");
|
||||
} else {
|
||||
curMessageList.push(curImage);
|
||||
}
|
||||
});
|
||||
curMessageList.push(
|
||||
message_str + this.config.seps[i % this.config.seps.length],
|
||||
);
|
||||
ret.push(curMessageList);
|
||||
}
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get prompt arrays with the first one as system.
|
||||
*
|
||||
* It is returned as an array of `string | Array<string | ImageURL>`, where each element of
|
||||
* the array represents the formatted message of a role/turn. If the message only contains text,
|
||||
* it will be a string that concatenates the role string, message, and separators. If the
|
||||
* message contains image(s), it will be an array of string and ImageURL in the order of which
|
||||
* they will be prefilled into the model. e.g. it can be something like
|
||||
* [
|
||||
* "<|system|>\nSome system prompt\n",
|
||||
* [
|
||||
* "<|user|>\n",
|
||||
* imageURL1,
|
||||
* "\n",
|
||||
* imageURL2,
|
||||
* "\n",
|
||||
* "Some user input<|end|>\n"
|
||||
* ],
|
||||
* ]
|
||||
*
|
||||
* @returns The prompt array.
|
||||
*/
|
||||
getPromptArray(
|
||||
config?: ChatConfig,
|
||||
): Array<string | Array<string | ImageURL>> {
|
||||
if (this.isTextCompletion) {
|
||||
throw new TextCompletionConversationError("getPromptArray");
|
||||
}
|
||||
return this.getPromptArrayInternal(true, 0, config);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the last round of prompt has not been fed as input.
|
||||
*
|
||||
* @note This function needs to be used with the assumption that
|
||||
* the caller call appendMessage then appendReplyHeader.
|
||||
*
|
||||
* @returns The prompt array.
|
||||
*/
|
||||
getPromptArrayLastRound(config?: ChatConfig) {
|
||||
if (this.isTextCompletion) {
|
||||
throw new TextCompletionConversationError("getPromptArrayLastRound");
|
||||
}
|
||||
if (this.messages.length < 3) {
|
||||
throw Error("needs to call getPromptArray for the first message");
|
||||
}
|
||||
return this.getPromptArrayInternal(false, this.messages.length - 2, config);
|
||||
}
|
||||
|
||||
/**
|
||||
* Return prompt in an array for non-conversation text completion.
|
||||
*/
|
||||
getPromptArrayTextCompletion(): Array<string> {
|
||||
if (!this.isTextCompletion || this.prompt === undefined) {
|
||||
throw new TextCompletionConversationExpectsPrompt();
|
||||
}
|
||||
return [this.prompt];
|
||||
}
|
||||
|
||||
/**
|
||||
* Resets all states for this.conversation.
|
||||
*/
|
||||
reset() {
|
||||
// Note: Update this whenever we introduce a new state to conversation.
|
||||
this.messages = [];
|
||||
this.override_system_message = undefined;
|
||||
this.function_string = "";
|
||||
this.use_function_calling = false;
|
||||
this.isTextCompletion = false;
|
||||
this.prompt = undefined;
|
||||
}
|
||||
|
||||
getStopStr(): string[] {
|
||||
// TODO(Charlie): Is this needed?
|
||||
// if (this.config.stop_str.length > 0) {
|
||||
// return this.config.stop_str;
|
||||
// }
|
||||
// return [this.config.seps[this.config.seps.length - 1]];
|
||||
return this.config.stop_str;
|
||||
}
|
||||
|
||||
getStopTokens() {
|
||||
return this.config.stop_token_ids;
|
||||
}
|
||||
|
||||
appendMessage(
|
||||
role: Role,
|
||||
message: string | Array<ChatCompletionContentPart>,
|
||||
role_name?: string,
|
||||
) {
|
||||
if (this.isTextCompletion) {
|
||||
throw new TextCompletionConversationError("appendMessage");
|
||||
}
|
||||
if (
|
||||
this.messages.length != 0 &&
|
||||
this.messages[this.messages.length - 1][2] == undefined
|
||||
) {
|
||||
throw Error("Have unfinished reply");
|
||||
}
|
||||
if (!(role in this.config.roles)) {
|
||||
throw Error("Role is not supported: " + role);
|
||||
}
|
||||
const role_name_str = role_name ? role_name : this.config.roles[role];
|
||||
this.messages.push([role, role_name_str, message]);
|
||||
}
|
||||
|
||||
appendReplyHeader(role: Role) {
|
||||
if (this.isTextCompletion) {
|
||||
throw new TextCompletionConversationError("appendReplyHeader");
|
||||
}
|
||||
if (!(role in this.config.roles)) {
|
||||
throw Error("Role is not supported: " + role);
|
||||
}
|
||||
this.messages.push([role, this.config.roles[role], undefined]);
|
||||
}
|
||||
|
||||
appendEmptyThinkingReplyHeader(role: Role, emptyThinkingBlockStr: string) {
|
||||
if (this.isTextCompletion) {
|
||||
throw new TextCompletionConversationError(
|
||||
"appendEmptyThinkingReplyHeader",
|
||||
);
|
||||
}
|
||||
this.isLastMessageEmptyThinkingReplyHeader = true;
|
||||
this.messages.push([role, this.config.roles[role], emptyThinkingBlockStr]);
|
||||
}
|
||||
|
||||
finishReply(message: string) {
|
||||
if (this.isTextCompletion) {
|
||||
throw new TextCompletionConversationError("finishReply");
|
||||
}
|
||||
if (this.messages.length == 0) {
|
||||
throw Error("Message error should not be 0");
|
||||
}
|
||||
if (
|
||||
this.messages[this.messages.length - 1][2] !== undefined &&
|
||||
// If the last message has an empty thinknig block, last message is expected
|
||||
// to be non-empty.
|
||||
this.isLastMessageEmptyThinkingReplyHeader === false
|
||||
) {
|
||||
throw Error("Already assigned");
|
||||
}
|
||||
this.messages[this.messages.length - 1][2] = message;
|
||||
this.isLastMessageEmptyThinkingReplyHeader = false;
|
||||
}
|
||||
}
|
||||
|
||||
export function getConversation(
|
||||
conv_template: ConvTemplateConfig,
|
||||
conv_config?: Partial<ConvTemplateConfig>,
|
||||
isTextCompletion = false,
|
||||
): Conversation {
|
||||
return new Conversation(
|
||||
{ ...conv_template, ...conv_config },
|
||||
isTextCompletion,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Compare the states of two conversation instances. Equality is defined as their getPromptArray()
|
||||
* should return the exact same things, which is determined by fields: messages, function_string,
|
||||
* use_function_calling, and override_system_message.
|
||||
*
|
||||
* @returns True if `convA` equals to `convB`
|
||||
* @note We assume convA and convB has the same `this.config`.
|
||||
*/
|
||||
export function compareConversationObject(
|
||||
convA: Conversation,
|
||||
convB: Conversation,
|
||||
): boolean {
|
||||
// NOTE: Update this function whenever a new state is introduced to `Conversation`.
|
||||
// Check the easy ones first
|
||||
if (
|
||||
convA.function_string !== convB.function_string ||
|
||||
convA.use_function_calling !== convB.use_function_calling ||
|
||||
convA.override_system_message !== convB.override_system_message ||
|
||||
convA.messages.length !== convB.messages.length ||
|
||||
convA.isTextCompletion !== convB.isTextCompletion
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Then check message
|
||||
if (convA.messages.length === 0 && convB.messages.length === 0) {
|
||||
// both are empty
|
||||
return true;
|
||||
}
|
||||
if (convA.messages.length !== convB.messages.length) {
|
||||
// different number of messages
|
||||
return false;
|
||||
}
|
||||
|
||||
const msgLen = convA.messages.length;
|
||||
const msgEntryLen = convA.messages[0].length; // always 3 for now
|
||||
for (let i = 0; i < msgLen; i++) {
|
||||
for (let j = 0; j < msgEntryLen; j++) {
|
||||
const entryA = convA.messages[i][j];
|
||||
const entryB = convB.messages[i][j];
|
||||
if (typeof entryA === "string" && typeof entryB === "string") {
|
||||
// Case 1: both are strings
|
||||
if (convA.messages[i][j] !== convB.messages[i][j]) {
|
||||
return false;
|
||||
}
|
||||
} else if (entryA === undefined && entryB === undefined) {
|
||||
// Case 2: both undefined
|
||||
continue;
|
||||
} else if (Array.isArray(entryA) && Array.isArray(entryB)) {
|
||||
// Case 3: both are ChatCompletionContentPart[]
|
||||
if (entryA.length !== entryB.length) {
|
||||
return false;
|
||||
}
|
||||
const numContentParts = entryA.length;
|
||||
for (let k = 0; k < numContentParts; k++) {
|
||||
const entryA_k = entryA[k];
|
||||
const entryB_k = entryB[k];
|
||||
if (entryA_k.type === "text" && entryB_k.type === "text") {
|
||||
// Case 3.1: both are text
|
||||
if (entryA_k.text !== entryB_k.text) {
|
||||
return false;
|
||||
}
|
||||
} else if (
|
||||
entryA_k.type === "image_url" &&
|
||||
entryB_k.type === "image_url"
|
||||
) {
|
||||
// Case 3.2: both are image_url
|
||||
if (
|
||||
entryA_k.image_url.url !== entryB_k.image_url.url ||
|
||||
entryA_k.image_url.detail !== entryB_k.image_url.detail
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
// Case 3.3: of different type
|
||||
return false;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Case 4: two entries are of different types
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get a new Conversation object based on the chat completion request.
|
||||
*
|
||||
* @param request The incoming ChatCompletionRequest
|
||||
* @param includeLastMsg Include last message, by default is false. Set to true for testing only.
|
||||
* @note By default, `request.messages[-1]` is not included as it would be treated as a normal
|
||||
* input to `prefill()`.
|
||||
*/
|
||||
export function getConversationFromChatCompletionRequest(
|
||||
request: ChatCompletionRequest,
|
||||
config: ChatConfig,
|
||||
includeLastMsg = false,
|
||||
): Conversation {
|
||||
// 0. Instantiate a new Conversation object
|
||||
const conversation = getConversation(
|
||||
config.conv_template,
|
||||
config.conv_config,
|
||||
);
|
||||
|
||||
// 1. Populate function-calling-related fields
|
||||
// TODO: either remove these or support gorilla-like function calling models.
|
||||
// These commented code was used to support gorilla, but we could not use grammar to
|
||||
// guarantee its output, nor make it conform to OpenAI's function calling output. Kept for now.
|
||||
// const functionCallUsage = this.getFunctionCallUsage(request);
|
||||
// conversation.function_string = functionCallUsage;
|
||||
// conversation.use_function_calling = functionCallUsage !== "";
|
||||
|
||||
// 2. Populate conversation.messages
|
||||
const input = request.messages;
|
||||
const lastId = input.length - 1;
|
||||
if (input[lastId].role !== "user" && input[lastId].role !== "tool") {
|
||||
throw new MessageOrderError(
|
||||
"The last message should be from the `user` or `tool`.",
|
||||
);
|
||||
}
|
||||
const iterEnd = includeLastMsg ? input.length : input.length - 1;
|
||||
for (let i = 0; i < iterEnd; i++) {
|
||||
const message: ChatCompletionMessageParam = input[i];
|
||||
if (message.role === "system") {
|
||||
if (i !== 0) {
|
||||
throw new SystemMessageOrderError();
|
||||
}
|
||||
conversation.override_system_message = message.content;
|
||||
} else if (message.role === "user") {
|
||||
conversation.appendMessage(Role.user, message.content, message.name);
|
||||
} else if (message.role === "assistant") {
|
||||
if (typeof message.content !== "string") {
|
||||
throw new ContentTypeError(message.role + "'s message");
|
||||
}
|
||||
conversation.appendMessage(Role.assistant, message.content, message.name);
|
||||
} else if (message.role === "tool") {
|
||||
conversation.appendMessage(Role.tool, message.content);
|
||||
} else {
|
||||
// Use `["role"]` instead of `.role` to suppress "Property does not exist on type 'never'"
|
||||
throw new UnsupportedRoleError(message["role"]);
|
||||
}
|
||||
}
|
||||
return conversation;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the function string based on the request.tools and request.tool_choice, raises erros if
|
||||
* encounter invalid request.
|
||||
*
|
||||
* @param request The chatCompletionRequest we are about to prefill for.
|
||||
* @returns The string used to set Conversation.function_string
|
||||
*/
|
||||
export function getFunctionCallUsage(request: ChatCompletionRequest): string {
|
||||
if (
|
||||
request.tools == undefined ||
|
||||
(typeof request.tool_choice == "string" && request.tool_choice == "none")
|
||||
) {
|
||||
return "";
|
||||
}
|
||||
if (
|
||||
typeof request.tool_choice == "string" &&
|
||||
request.tool_choice !== "auto"
|
||||
) {
|
||||
throw new InvalidToolChoiceError(request.tool_choice);
|
||||
}
|
||||
if (
|
||||
typeof request.tool_choice !== "string" &&
|
||||
request.tool_choice?.type !== "function"
|
||||
) {
|
||||
throw new UnsupportedToolChoiceTypeError();
|
||||
}
|
||||
|
||||
const singleFunctionToCall =
|
||||
typeof request.tool_choice !== "string" &&
|
||||
request.tool_choice?.function?.name;
|
||||
if (singleFunctionToCall) {
|
||||
for (const f of request.tools) {
|
||||
if (singleFunctionToCall == f.function.name) {
|
||||
return JSON.stringify([f.function]);
|
||||
}
|
||||
}
|
||||
throw new FunctionNotFoundError(singleFunctionToCall);
|
||||
}
|
||||
|
||||
const function_list = [];
|
||||
for (const f of request.tools) {
|
||||
if (f.type !== "function") {
|
||||
throw new UnsupportedToolTypeError();
|
||||
}
|
||||
function_list.push(f.function);
|
||||
}
|
||||
return JSON.stringify(function_list);
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
import * as tvmjs from "@mlc-ai/web-runtime";
|
||||
import log from "loglevel";
|
||||
import { Tokenizer } from "@mlc-ai/web-tokenizers";
|
||||
import { ChatConfig } from "./config";
|
||||
import {
|
||||
EmbeddingChunkingUnsupportedError,
|
||||
EmbeddingExceedContextWindowSizeError,
|
||||
EmbeddingInputEmptyError,
|
||||
EmbeddingSlidingWindowError,
|
||||
MinValueError,
|
||||
} from "./error";
|
||||
|
||||
export class EmbeddingPipeline {
|
||||
private config: ChatConfig;
|
||||
private tokenizer: Tokenizer;
|
||||
|
||||
// TVM functions
|
||||
private tvm: tvmjs.Instance;
|
||||
private device: tvmjs.DLDevice;
|
||||
private vm: tvmjs.VirtualMachine;
|
||||
private prefill: tvmjs.PackedFunc;
|
||||
private params: tvmjs.TVMObject;
|
||||
|
||||
// metadata
|
||||
private contextWindowSize = -1;
|
||||
private prefillChunkSize = -1;
|
||||
private maxBatchSize = -1;
|
||||
|
||||
// performance
|
||||
private curRoundEmbedTotalTokens = 0; // excludes padded tokens for batching
|
||||
private curRoundEmbedTotalTime = 0;
|
||||
|
||||
constructor(tvm: tvmjs.Instance, tokenizer: Tokenizer, config: ChatConfig) {
|
||||
// 0. Setting attributes
|
||||
this.tvm = tvm;
|
||||
this.tokenizer = tokenizer;
|
||||
this.config = config;
|
||||
this.device = this.tvm.webgpu();
|
||||
|
||||
// 1. Create VM and get the core functions
|
||||
tvm.beginScope();
|
||||
this.vm = this.tvm.detachFromCurrentScope(
|
||||
this.tvm.createVirtualMachine(this.device),
|
||||
);
|
||||
this.prefill = this.tvm.detachFromCurrentScope(
|
||||
this.vm.getFunction("prefill"),
|
||||
);
|
||||
|
||||
// 2. Get json stored in the vm's metadata function
|
||||
const fgetMetadata = this.vm.getFunction("_metadata");
|
||||
const ret_value = fgetMetadata();
|
||||
const metadataStr = ret_value.toString();
|
||||
const metadata = JSON.parse(metadataStr);
|
||||
|
||||
// 3. Load parameters by name
|
||||
const paramNames: string[] = [];
|
||||
metadata.params.forEach((param: any) => {
|
||||
paramNames.push(param.name);
|
||||
});
|
||||
this.params = this.tvm.detachFromCurrentScope(
|
||||
this.tvm.getParamsFromCacheByName(paramNames),
|
||||
);
|
||||
|
||||
// 4. Read in compilation configurations from metadata
|
||||
// We use context window size max batch size to check validity of the model
|
||||
// We assume prefillChunkSize is the same as contextWindowSize for embedding model for now
|
||||
this.maxBatchSize = metadata.max_batch_size;
|
||||
this.contextWindowSize = this.config.context_window_size;
|
||||
this.prefillChunkSize = metadata.prefill_chunk_size;
|
||||
log.info("Using maxBatchSize: ", this.maxBatchSize);
|
||||
log.info("Using contextWindowSize: ", this.contextWindowSize);
|
||||
log.info("Using prefillChunkSize: ", this.prefillChunkSize);
|
||||
|
||||
if (this.config.sliding_window_size !== -1) {
|
||||
throw new EmbeddingSlidingWindowError(this.config.sliding_window_size);
|
||||
}
|
||||
if (this.maxBatchSize <= 0) {
|
||||
throw new MinValueError("maxBatchSize", 0);
|
||||
}
|
||||
if (this.contextWindowSize <= 0) {
|
||||
throw new MinValueError("contextWindowSize", 0);
|
||||
}
|
||||
if (this.prefillChunkSize <= 0) {
|
||||
throw new MinValueError("prefillChunkSize", 0);
|
||||
}
|
||||
if (this.prefillChunkSize !== this.contextWindowSize) {
|
||||
throw new EmbeddingChunkingUnsupportedError(
|
||||
this.contextWindowSize,
|
||||
this.prefillChunkSize,
|
||||
);
|
||||
}
|
||||
tvm.endScope();
|
||||
}
|
||||
|
||||
async embedStep(
|
||||
input: string | Array<string> | Array<number> | Array<Array<number>>,
|
||||
): Promise<Array<Array<number>>> {
|
||||
// 0. Reset performance metrics
|
||||
this.curRoundEmbedTotalTokens = 0;
|
||||
this.curRoundEmbedTotalTime = 0;
|
||||
let totalNumTokens = 0;
|
||||
const embedStart = performance.now();
|
||||
let tokenizedInputs: Array<Array<number>> = [];
|
||||
const tempInputs: Array<number> = [];
|
||||
// 1. Convert all possible input types to Array<Array<number>>, tokenize if not already
|
||||
// Cannot use input.every to match type, which leads to TS compilation error
|
||||
// https://github.com/microsoft/TypeScript/issues/33591
|
||||
if (input.length === 0) {
|
||||
throw new EmbeddingInputEmptyError();
|
||||
}
|
||||
if (typeof input === "string") {
|
||||
// string
|
||||
tokenizedInputs = [Array.from(this.tokenizer.encode(input))];
|
||||
} else {
|
||||
for (let i = 0; i < input.length; i++) {
|
||||
const curInput = input[i];
|
||||
if (Array.isArray(curInput)) {
|
||||
// Array<Array<number>>
|
||||
tokenizedInputs.push(curInput);
|
||||
} else if (typeof curInput === "string") {
|
||||
// Array<string>
|
||||
tokenizedInputs.push(Array.from(this.tokenizer.encode(curInput)));
|
||||
} else {
|
||||
// Array<number>
|
||||
tempInputs.push(curInput);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (tempInputs.length > 0) {
|
||||
tokenizedInputs.push(tempInputs);
|
||||
}
|
||||
|
||||
// 2. Check each input is not larger than the context window size
|
||||
// TODO: tokenizer.encode seems to implicitly truncates to contextWindowSize, confirm behavior
|
||||
// and decide whether to warn user
|
||||
for (let i = 0; i < tokenizedInputs.length; i++) {
|
||||
const curInputSize = tokenizedInputs[i].length;
|
||||
totalNumTokens += curInputSize;
|
||||
if (curInputSize > this.contextWindowSize) {
|
||||
throw new EmbeddingExceedContextWindowSizeError(
|
||||
this.contextWindowSize,
|
||||
curInputSize,
|
||||
);
|
||||
}
|
||||
}
|
||||
if (tokenizedInputs.length === 0) {
|
||||
throw new Error("InternalError: batch size is zero.");
|
||||
}
|
||||
|
||||
// 3. Forward each batch
|
||||
const batchSize = tokenizedInputs.length;
|
||||
const result: Array<Array<number>> = [];
|
||||
for (let begin = 0; begin < batchSize; begin += this.maxBatchSize) {
|
||||
this.tvm.beginScope();
|
||||
// 3.1 Get current batch
|
||||
const end = Math.min(batchSize, begin + this.maxBatchSize);
|
||||
const curBatch: Array<Array<number>> = tokenizedInputs.slice(begin, end);
|
||||
const curBatchSize = curBatch.length;
|
||||
// 3.2 Max input size of current batch
|
||||
let maxInputSize = 0;
|
||||
for (let i = 0; i < curBatchSize; i++) {
|
||||
const curInputSize = curBatch[i].length;
|
||||
if (curInputSize > maxInputSize) {
|
||||
maxInputSize = curInputSize;
|
||||
}
|
||||
}
|
||||
// 3.3 Create inputs and attention mask
|
||||
// Padded with zeros and flattened, of size curBatchSize * maxInputSize
|
||||
const curBatchPaddedFlatten: Array<number> = [];
|
||||
// 1 for non-pad, 0 otherwise, also of size curBatchSize * maxInputSize
|
||||
const curAttnMask: Array<number> = [];
|
||||
const flattenedInputSize = curBatchSize * maxInputSize;
|
||||
for (let i = 0; i < curBatchSize; i++) {
|
||||
const padding = Array(maxInputSize - curBatch[i].length).fill(0);
|
||||
const ones = Array(curBatch[i].length).fill(1);
|
||||
curBatchPaddedFlatten.push(...curBatch[i]);
|
||||
curAttnMask.push(...ones);
|
||||
curBatchPaddedFlatten.push(...padding);
|
||||
curAttnMask.push(...padding);
|
||||
}
|
||||
if (
|
||||
curBatchPaddedFlatten.length !== flattenedInputSize ||
|
||||
curAttnMask.length !== flattenedInputSize
|
||||
) {
|
||||
throw new Error(
|
||||
`InternalError: Expect input array to be ${flattenedInputSize}, ` +
|
||||
`but got ${curBatchPaddedFlatten.length}`,
|
||||
);
|
||||
}
|
||||
// 3.4 Convert inputs and attention mask to tvm ndarray on GPU, of shape (curBatchSize, maxInputSize)
|
||||
let inputNDArray = this.tvm.empty(
|
||||
[flattenedInputSize],
|
||||
"int32",
|
||||
this.device,
|
||||
);
|
||||
inputNDArray.copyFrom(curBatchPaddedFlatten);
|
||||
inputNDArray = inputNDArray.view([curBatchSize, maxInputSize]);
|
||||
let maskNDArray = this.tvm.empty(
|
||||
[flattenedInputSize],
|
||||
"int32",
|
||||
this.device,
|
||||
);
|
||||
maskNDArray.copyFrom(curAttnMask);
|
||||
maskNDArray = maskNDArray.view([curBatchSize, maxInputSize]);
|
||||
|
||||
// 3.5 Actual forwarding on GPU, logits of shape (curBatchSize, maxInputSize, hidden_size)
|
||||
const logitsCurBatchOnGPU: tvmjs.Tensor = this.prefill(
|
||||
inputNDArray,
|
||||
maskNDArray,
|
||||
this.params,
|
||||
);
|
||||
await this.device.sync();
|
||||
|
||||
// 3.6 Copy logits to CPU, flatten to curBatchSize * maxInputSize * hidden_size
|
||||
const hidden_size = logitsCurBatchOnGPU.shape[2];
|
||||
let logitsCurBatchOnCPU: tvmjs.Tensor = this.tvm.empty(
|
||||
logitsCurBatchOnGPU.shape,
|
||||
logitsCurBatchOnGPU.dtype,
|
||||
this.tvm.cpu(),
|
||||
);
|
||||
logitsCurBatchOnCPU.copyFrom(logitsCurBatchOnGPU);
|
||||
logitsCurBatchOnCPU = logitsCurBatchOnCPU.view([
|
||||
curBatchSize * maxInputSize * hidden_size,
|
||||
]);
|
||||
await this.device.sync();
|
||||
const logitsCurBatchOnCPUArray: Float32Array = <Float32Array>(
|
||||
logitsCurBatchOnCPU.toArray()
|
||||
);
|
||||
|
||||
// 3.7 Update final result. For each sentence, get [0,:], i.e. only the first token's output
|
||||
// That is, we are doing result.push(logits[:,0,:]) here.
|
||||
// TODO: check if all models only use [0,:]. If it is snowflake-specific, need to specify
|
||||
// this in mlc-chat-config.json
|
||||
for (let i = 0; i < curBatchSize; i++) {
|
||||
const b = i * maxInputSize * hidden_size;
|
||||
const e = b + hidden_size;
|
||||
result.push(Array.from(logitsCurBatchOnCPUArray.slice(b, e)));
|
||||
}
|
||||
this.tvm.endScope();
|
||||
}
|
||||
if (result.length !== batchSize) {
|
||||
throw new Error(`
|
||||
InternalError: expect result.length to be ${batchSize}, but got ${result.length}`);
|
||||
}
|
||||
const embedEnd = performance.now();
|
||||
this.curRoundEmbedTotalTokens = totalNumTokens;
|
||||
this.curRoundEmbedTotalTime = (embedEnd - embedStart) / 1e3;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
dispose() {
|
||||
this.params.dispose();
|
||||
this.prefill.dispose();
|
||||
this.vm.dispose();
|
||||
this.tvm.dispose();
|
||||
this.tokenizer.dispose();
|
||||
}
|
||||
|
||||
/**
|
||||
* Synchronize the device.
|
||||
*/
|
||||
async sync(): Promise<void> {
|
||||
// Is it equivalent to this.tvm.sync()?
|
||||
await this.device.sync();
|
||||
}
|
||||
|
||||
async asyncLoadWebGPUPipelines() {
|
||||
await this.tvm.asyncLoadWebGPUPipelines(this.vm.getInternalModule());
|
||||
}
|
||||
|
||||
// Performance APIs below
|
||||
|
||||
/**
|
||||
* Get the time it took the last `embedStep()` in seconds.
|
||||
*/
|
||||
getCurRoundEmbedTotalTime(): number {
|
||||
return this.curRoundEmbedTotalTime;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the number of tokens embedded in the last `embedStep()`. This excludes the padded tokens.
|
||||
*/
|
||||
getCurRoundEmbedTotalTokens(): number {
|
||||
return this.curRoundEmbedTotalTokens;
|
||||
}
|
||||
|
||||
/**
|
||||
* @returns Prefill tokens per second, starting from the last prefill performed.
|
||||
*/
|
||||
getCurRoundEmbedTokensPerSec(): number {
|
||||
return this.curRoundEmbedTotalTokens / this.curRoundEmbedTotalTime;
|
||||
}
|
||||
}
|
||||
+1432
File diff suppressed because it is too large
Load Diff
+629
@@ -0,0 +1,629 @@
|
||||
export class ModelNotFoundError extends Error {
|
||||
constructor(modelId: string) {
|
||||
super(
|
||||
`Cannot find model record in appConfig for ${modelId}. Please check if the model ID is correct and included in the model_list configuration.`,
|
||||
);
|
||||
this.name = "ModelNotFoundError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ConfigValueError extends Error {
|
||||
constructor(message: string) {
|
||||
super(message);
|
||||
this.name = "ConfigValueError";
|
||||
}
|
||||
}
|
||||
|
||||
export class MinValueError extends ConfigValueError {
|
||||
constructor(paramName: string, minValue: number) {
|
||||
super(`Make sure \`${paramName}\` > ${minValue}.`);
|
||||
this.name = "MinValueError";
|
||||
}
|
||||
}
|
||||
|
||||
export class RangeError extends ConfigValueError {
|
||||
constructor(
|
||||
paramName: string,
|
||||
minValue: number,
|
||||
maxValue: number,
|
||||
additionalMessage?: string,
|
||||
) {
|
||||
super(
|
||||
`Make sure ${minValue} < ${paramName} <= ${maxValue}.${additionalMessage ? " " + additionalMessage : ""}`,
|
||||
);
|
||||
this.name = "RangeError";
|
||||
}
|
||||
}
|
||||
|
||||
export class NonNegativeError extends ConfigValueError {
|
||||
constructor(paramName: string) {
|
||||
super(`Make sure ${paramName} >= 0.`);
|
||||
this.name = "NonNegativeError";
|
||||
}
|
||||
}
|
||||
|
||||
export class InvalidNumberStringError extends ConfigValueError {
|
||||
constructor(paramName: string, actualValue?: string) {
|
||||
super(
|
||||
`Make sure ${paramName} to be number represented in string.${actualValue ? " Got " + actualValue : ""}`,
|
||||
);
|
||||
this.name = "InvalidNumberStringError";
|
||||
}
|
||||
}
|
||||
|
||||
export class DependencyError extends ConfigValueError {
|
||||
constructor(
|
||||
dependentParam: string,
|
||||
requiredParam: string,
|
||||
requiredValue: any,
|
||||
) {
|
||||
super(
|
||||
`${dependentParam} requires ${requiredParam} to be ${requiredValue}.`,
|
||||
);
|
||||
this.name = "DependencyError";
|
||||
}
|
||||
}
|
||||
|
||||
export class WebGPUNotAvailableError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"WebGPU is not supported in your current environment, but it is necessary to run the WebLLM engine. " +
|
||||
"Please make sure that your browser supports WebGPU and that it is enabled in your browser settings. " +
|
||||
"You can also consult your browser's compatibility chart to see if it supports WebGPU. " +
|
||||
"For more information about WebGPU support in your browser, visit https://webgpureport.org/",
|
||||
);
|
||||
this.name = "WebGPUNotAvailableError";
|
||||
}
|
||||
}
|
||||
|
||||
export class WebGPUNotFoundError extends Error {
|
||||
constructor() {
|
||||
super("Cannot find WebGPU in the environment");
|
||||
this.name = "WebGPUNotFoundError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ModelNotLoadedError extends Error {
|
||||
constructor(requestName: string) {
|
||||
super(
|
||||
`Model not loaded before trying to complete ${requestName}. Please ensure you have called ` +
|
||||
`MLCEngine.reload(model) to load the model before initiating APIs, ` +
|
||||
`or initialize your engine using CreateMLCEngine() with a valid model configuration.`,
|
||||
);
|
||||
this.name = "ModelNotLoadedError";
|
||||
}
|
||||
}
|
||||
|
||||
export class WorkerEngineModelNotLoadedError extends Error {
|
||||
constructor(engineName: string) {
|
||||
super(
|
||||
`${engineName} is not loaded with a model. Did you call \`engine.reload()\`?`,
|
||||
);
|
||||
this.name = "WorkerEngineModelNotLoadedError";
|
||||
}
|
||||
}
|
||||
|
||||
export class MessageOrderError extends Error {
|
||||
constructor(message: string) {
|
||||
super(message);
|
||||
this.name = "MessageOrderError";
|
||||
}
|
||||
}
|
||||
|
||||
export class SystemMessageOrderError extends Error {
|
||||
constructor() {
|
||||
super("System prompt should always be the first message in `messages`.");
|
||||
this.name = "SystemMessageOrderError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ContentTypeError extends Error {
|
||||
constructor(name: string) {
|
||||
super(`${name} should have string content.`);
|
||||
this.name = "ContentTypeError";
|
||||
}
|
||||
}
|
||||
|
||||
export class UnsupportedRoleError extends Error {
|
||||
constructor(role: string) {
|
||||
super(`Unsupported role of message: ${role}`);
|
||||
this.name = "UnsupportedRoleError";
|
||||
}
|
||||
}
|
||||
|
||||
export class UserMessageContentErrorForNonVLM extends Error {
|
||||
constructor(modelId: string, modelType: string, content: any) {
|
||||
super(
|
||||
`The model loaded is not of type ModelType.VLM (vision-language model). ` +
|
||||
`Therefore, user message only supports string content, but received: ${content}\n` +
|
||||
`Loaded modelId: ${modelId}, modelType: ${modelType}`,
|
||||
);
|
||||
this.name = "UserMessageContentErrorForNonVLM";
|
||||
}
|
||||
}
|
||||
|
||||
export class PrefillChunkSizeSmallerThanImageError extends Error {
|
||||
constructor(prefillChunkSize: number, imageEmbedSize: number) {
|
||||
super(
|
||||
`prefillChunkSize needs to be greater than imageEmbedSize because a single image's ` +
|
||||
`prefill cannot be chunked. Got prefillChunkSize: ` +
|
||||
`${prefillChunkSize}, imageEmbedSize: ${imageEmbedSize}`,
|
||||
);
|
||||
this.name = "PrefillChunkSizeSmallerThanImageError";
|
||||
}
|
||||
}
|
||||
|
||||
export class CannotFindImageEmbedError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
`Received image input but model does not have kernel image_embed. ` +
|
||||
`Make sure to only pass in image to a vision model.`,
|
||||
);
|
||||
this.name = "CannotFindImageEmbedError";
|
||||
}
|
||||
}
|
||||
|
||||
export class UnsupportedDetailError extends Error {
|
||||
constructor(detail: string) {
|
||||
super(
|
||||
`Currently do not support field image_url.detail, but received: ${detail}`,
|
||||
);
|
||||
this.name = "UnsupportedDetailError";
|
||||
}
|
||||
}
|
||||
|
||||
export class UnsupportedImageURLError extends Error {
|
||||
constructor(url: string) {
|
||||
super(
|
||||
`image_url.url should start with "data:image" for base64, or with "http", but got: ${url}`,
|
||||
);
|
||||
this.name = "UnsupportedImageURLError";
|
||||
}
|
||||
}
|
||||
|
||||
export class MultipleTextContentError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
`Each message can have at most one text contentPart, but received more than 1.`,
|
||||
);
|
||||
this.name = "MultipleTextContentError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ToolCallOutputParseError extends Error {
|
||||
constructor(outputMessage: string, error: Error) {
|
||||
super(
|
||||
`Internal error: error encountered when parsing outputMessage for function ` +
|
||||
`calling. Got outputMessage: ${outputMessage}\nGot error: ${error}`,
|
||||
);
|
||||
this.name = "ToolCallOutputParseError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ToolCallOutputInvalidTypeError extends Error {
|
||||
constructor(expectedType: string) {
|
||||
super(
|
||||
`Internal error: expect output of function calling to be an ${expectedType}`,
|
||||
);
|
||||
this.name = "ToolCallOutputInvalidTypeError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ToolCallOutputMissingFieldsError extends Error {
|
||||
constructor(missingFields: string[], object: any) {
|
||||
super(
|
||||
`Expect generated tool call to have fields ${missingFields.map((field) => `"\`${field}\`"`).join(", ")}, but got object: ${JSON.stringify(object)}`,
|
||||
);
|
||||
this.name = "JSONFieldError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ConfigurationNotInitializedError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Configuration not initialized. Ensure you have called `reload()` function first.",
|
||||
);
|
||||
this.name = "ConfigurationNotInitializedError";
|
||||
}
|
||||
}
|
||||
|
||||
export class MissingModelWasmError extends Error {
|
||||
constructor(modelId: string) {
|
||||
super(
|
||||
`Missing \`model_lib\` for the model with ID "${modelId}". Please ensure that \`model_lib\` is provided in \`model_list\` for each model. This URL is essential for downloading the WASM library necessary to run the model.`,
|
||||
);
|
||||
this.name = "MissingModelError";
|
||||
}
|
||||
}
|
||||
|
||||
export class FeatureSupportError extends Error {
|
||||
constructor(feature: string) {
|
||||
super(
|
||||
`This model requires feature ${feature}, which is not yet supported by this browser.`,
|
||||
);
|
||||
this.name = "FeatureSupportError";
|
||||
}
|
||||
}
|
||||
|
||||
export class UnsupportedFieldsError extends Error {
|
||||
constructor(unsupportedFields: string[], targetClass: string) {
|
||||
super(
|
||||
`The following fields in ${targetClass} are not yet supported: \n` +
|
||||
unsupportedFields.join(", "),
|
||||
);
|
||||
this.name = "UnsupportedFieldsError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ShaderF16SupportError extends FeatureSupportError {
|
||||
constructor() {
|
||||
super(
|
||||
"This model requires WebGPU extension shader-f16, which is not enabled in this browser. " +
|
||||
'You can try to launch Chrome Canary in command line with flag "--enable-dawn-features=allow_unsafe_apis".',
|
||||
);
|
||||
this.name = "ShaderF16SupportError";
|
||||
}
|
||||
}
|
||||
export class DeviceLostError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"The WebGPU device was lost while loading the model. This issue often occurs due to running out of memory (OOM). To resolve this, try reloading with a model that has fewer parameters or uses a smaller context length.",
|
||||
);
|
||||
this.name = "DeviceLostError";
|
||||
}
|
||||
}
|
||||
|
||||
export class InvalidToolChoiceError extends Error {
|
||||
constructor(toolChoice: string) {
|
||||
super(
|
||||
`Invalid tool choice value: '${toolChoice}'. Please check your input and try again.`,
|
||||
);
|
||||
this.name = "InvalidToolChoiceError";
|
||||
}
|
||||
}
|
||||
|
||||
export class UnsupportedToolChoiceTypeError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Unsupported tool choice type. Only tool choices of type 'function' are supported.",
|
||||
);
|
||||
this.name = "UnsupportedToolChoiceTypeError";
|
||||
}
|
||||
}
|
||||
|
||||
export class FunctionNotFoundError extends Error {
|
||||
constructor(functionName: string) {
|
||||
super(
|
||||
`The tool choice function ${functionName} is not found in the tools list`,
|
||||
);
|
||||
this.name = "FunctionNotFoundError";
|
||||
}
|
||||
}
|
||||
|
||||
export class UnsupportedToolTypeError extends Error {
|
||||
constructor() {
|
||||
super("Only 'function' tool type is supported");
|
||||
this.name = "UnsupportedToolTypeError";
|
||||
}
|
||||
}
|
||||
export class EngineNotLoadedError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Engine not yet loaded with model. Ensure you initialize the chat module by calling `engine.reload()` first.",
|
||||
);
|
||||
this.name = "EngineNotLoadedError";
|
||||
}
|
||||
}
|
||||
export class UnsupportedTokenizerFilesError extends Error {
|
||||
constructor(files: string[]) {
|
||||
super(`Cannot handle tokenizer files ${files}`);
|
||||
this.name = "UnsupportedTokenizerFilesError";
|
||||
}
|
||||
}
|
||||
|
||||
export class WindowSizeConfigurationError extends Error {
|
||||
constructor(contextWindowSize: number, slidingWindowSize: number) {
|
||||
super(
|
||||
`Only one of context_window_size and sliding_window_size can be positive. Got: ` +
|
||||
`context_window_size: ${contextWindowSize}, sliding_window_size: ${slidingWindowSize}\n` +
|
||||
`Consider modifying ModelRecord.overrides to set one of them to -1.`,
|
||||
);
|
||||
this.name = "WindowSizeConfigurationError";
|
||||
}
|
||||
}
|
||||
|
||||
export class AttentionSinkSizeError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Need to specify non-negative attention_sink_size if using sliding window. " +
|
||||
"Consider modifying ModelRecord.overrides. " +
|
||||
"Use `attention_sink_size=0` for default sliding window.",
|
||||
);
|
||||
this.name = "AttentionSinkSizeError";
|
||||
}
|
||||
}
|
||||
|
||||
export class WindowSizeSpecificationError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Need to specify either sliding_window_size or max_window_size.\n" +
|
||||
"Consider modifying ModelRecord.overrides to set one of them to positive.",
|
||||
);
|
||||
this.name = "WindowSizeSpecificationError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ContextWindowSizeExceededError extends Error {
|
||||
constructor(numPromptTokens: number, contextWindowSize: number) {
|
||||
super(
|
||||
`Prompt tokens exceed context window size: number of prompt tokens: ${numPromptTokens}; ` +
|
||||
`context window size: ${contextWindowSize}\nConsider shortening the prompt, or increase ` +
|
||||
"`context_window_size`, or using sliding window via `sliding_window_size`.",
|
||||
);
|
||||
this.name = "ContextWindowSizeExceededError";
|
||||
}
|
||||
}
|
||||
|
||||
export class NonWorkerEnvironmentError extends Error {
|
||||
constructor(className: string) {
|
||||
super(`${className} must be created in the service worker script.`);
|
||||
this.name = "NonWorkerEnvironmentError";
|
||||
}
|
||||
}
|
||||
|
||||
export class NoServiceWorkerAPIError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Service worker API is not available in your browser. Please ensure that your browser supports service workers and that you are using a secure context (HTTPS). " +
|
||||
"Check the browser compatibility and ensure that service workers are not disabled in your browser settings.",
|
||||
);
|
||||
this.name = "NoServiceWorkerAPIError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ServiceWorkerInitializationError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Service worker failed to initialize. This could be due to a failure in the service worker registration process or because the service worker is not active. " +
|
||||
"Please refresh the page to retry initializing the service worker.",
|
||||
);
|
||||
this.name = "ServiceWorkerInitializationError";
|
||||
}
|
||||
}
|
||||
|
||||
export class StreamingCountError extends Error {
|
||||
constructor() {
|
||||
super("When streaming, `n` cannot be > 1.");
|
||||
this.name = "StreamingCountError";
|
||||
}
|
||||
}
|
||||
|
||||
export class SeedTypeError extends Error {
|
||||
constructor(seed: any) {
|
||||
super("`seed` should be an integer, but got " + seed);
|
||||
this.name = "SeedTypeError";
|
||||
}
|
||||
}
|
||||
export class InvalidResponseFormatError extends Error {
|
||||
constructor() {
|
||||
super("JSON schema is only supported with `json_object` response format.");
|
||||
this.name = "InvalidResponseFormatError";
|
||||
}
|
||||
}
|
||||
|
||||
export class InvalidResponseFormatGrammarError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"When ResponseFormat.type is `grammar`, ResponseFormat.grammar needs to be specified.\n" +
|
||||
"When ResponseFormat.grammar is specified, ResponseFormat.type needs to be grammar.",
|
||||
);
|
||||
this.name = "InvalidResponseFormatGrammarError";
|
||||
}
|
||||
}
|
||||
|
||||
export class InvalidResponseFormatStructuralTagError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"When ResponseFormat.type is `structural_tag`, ResponseFormat.structural_tag needs to be specified.\n" +
|
||||
"When ResponseFormat.structural_tag is specified, ResponseFormat.type needs to be structural_tag.",
|
||||
);
|
||||
this.name = "InvalidResponseFormatStructuralTagError";
|
||||
}
|
||||
}
|
||||
|
||||
export class CustomResponseFormatError extends Error {
|
||||
constructor(currentFormat: any) {
|
||||
super(
|
||||
"When using Hermes-2-Pro function calling via ChatCompletionRequest.tools, " +
|
||||
"cannot specify customized response_format. We will set it for you internally. Currently " +
|
||||
"set to: " +
|
||||
JSON.stringify(currentFormat),
|
||||
);
|
||||
this.name = "CustomResponseFormatError";
|
||||
}
|
||||
}
|
||||
export class UnsupportedModelIdError extends Error {
|
||||
constructor(currentModelId: string, supportedModelIds: string[]) {
|
||||
super(
|
||||
`${currentModelId} is not supported for ChatCompletionRequest.tools. Currently, models ` +
|
||||
`that support function calling are: ${supportedModelIds.join(", ")}`,
|
||||
);
|
||||
this.name = "UnsupportedModelIdError";
|
||||
}
|
||||
}
|
||||
export class CustomSystemPromptError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"When using Hermes-2-Pro function calling via ChatCompletionRequest.tools, cannot specify customized system prompt.",
|
||||
);
|
||||
this.name = "CustomSystemPromptError";
|
||||
}
|
||||
}
|
||||
|
||||
export class InvalidStreamOptionsError extends Error {
|
||||
constructor() {
|
||||
super("Only specify stream_options when stream=True.");
|
||||
this.name = "InvalidStreamOptionsError";
|
||||
}
|
||||
}
|
||||
export class UnknownMessageKindError extends Error {
|
||||
constructor(msgKind: string, msgContent: any) {
|
||||
super(`Unknown message kind, msg: [${msgKind}] ${msgContent}`);
|
||||
this.name = "UnknownMessageKindError";
|
||||
}
|
||||
}
|
||||
|
||||
export class TextCompletionExpectsKVEmptyError extends Error {
|
||||
constructor() {
|
||||
super("Non-chat text completion API expects KVCache to be empty.");
|
||||
this.name = "TextCompletionExpectsKVEmptyError";
|
||||
}
|
||||
}
|
||||
|
||||
export class TextCompletionConversationExpectsPrompt extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Non-chat text completion API expects isTextCompletion is true, and prompt is defined.",
|
||||
);
|
||||
this.name = "TextCompletionConversationExpectsPrompt";
|
||||
}
|
||||
}
|
||||
|
||||
export class TextCompletionConversationError extends Error {
|
||||
constructor(funcName: string) {
|
||||
super(`Non-chat text completion API cannot call ${funcName}.`);
|
||||
this.name = "TextCompletionConversationError";
|
||||
}
|
||||
}
|
||||
|
||||
export class EmbeddingUnsupportedEncodingFormatError extends Error {
|
||||
constructor() {
|
||||
super("Embedding in base64 format is currently not supported.");
|
||||
this.name = "EmbeddingUnsupportedEncodingFormatError";
|
||||
}
|
||||
}
|
||||
|
||||
export class EmbeddingUnsupportedModelError extends Error {
|
||||
constructor(currentModel: string) {
|
||||
super(
|
||||
`Trying to run embeddings.create() with ${currentModel}, which does not have ` +
|
||||
`ModelRecord.model_type === ModelType.embedding in the model record. ` +
|
||||
`Either make sure an embedding model is loaded, or specify the model type in ModelRecord.`,
|
||||
);
|
||||
this.name = "EmbeddingUnsupportedModelError";
|
||||
}
|
||||
}
|
||||
|
||||
export class EmbeddingSlidingWindowError extends Error {
|
||||
constructor(sliding_window_size: number) {
|
||||
super(
|
||||
`Embedding should not use sliding window. However, ` +
|
||||
`sliding_window_size=${sliding_window_size} is specified in the chat config.`,
|
||||
);
|
||||
this.name = "EmbeddingSlidingWindowError";
|
||||
}
|
||||
}
|
||||
|
||||
export class EmbeddingChunkingUnsupportedError extends Error {
|
||||
constructor(contextWindowSize: number, prefillChunkSize: number) {
|
||||
super(
|
||||
`Embedding currently does not support chunking. Make sure ` +
|
||||
`contextWindowSize === prefillChunkSize. Got contextWindowSize=${contextWindowSize}, ` +
|
||||
`prefillChunkSize=${prefillChunkSize} instead.`,
|
||||
);
|
||||
this.name = "EmbeddingChunkingUnsupportedError";
|
||||
}
|
||||
}
|
||||
|
||||
export class EmbeddingExceedContextWindowSizeError extends Error {
|
||||
constructor(contextWindowSize: number, receivedSize: number) {
|
||||
super(
|
||||
`The embedding model you are using only supports up to ${contextWindowSize} context size.` +
|
||||
`However, an input in the batch has size ${receivedSize}.`,
|
||||
);
|
||||
this.name = "EmbeddingExceedContextWindowSizeError";
|
||||
}
|
||||
}
|
||||
|
||||
export class EmbeddingInputEmptyError extends Error {
|
||||
constructor() {
|
||||
super("Embedding input cannot be empty string or empty token array.");
|
||||
this.name = "EmbeddingInputEmptyError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ReloadArgumentSizeUnmatchedError extends Error {
|
||||
constructor(numModelId: number, numChatOpts: number) {
|
||||
super(
|
||||
`Expect chatOpts, if specified, to match the size of modelId. However, got ` +
|
||||
`${numModelId} modelId, but ${numChatOpts} chatOpts.`,
|
||||
);
|
||||
this.name = "ReloadArgumentSizeUnmatchedError";
|
||||
}
|
||||
}
|
||||
|
||||
export class UnclearModelToUseError extends Error {
|
||||
constructor(loadedModels: string[], requestName: string) {
|
||||
super(
|
||||
`Multiple models are loaded in engine. Please specify the model in ${requestName}.\n` +
|
||||
`Currently loaded models are:\n${loadedModels}`,
|
||||
);
|
||||
this.name = "UnclearModelToUseError";
|
||||
}
|
||||
}
|
||||
|
||||
export class SpecifiedModelNotFoundError extends Error {
|
||||
constructor(
|
||||
loadedModels: string[],
|
||||
requestedModelId: string,
|
||||
requestName: string,
|
||||
) {
|
||||
super(
|
||||
`Specified model ${requestedModelId} for ${requestName} is not found in loaded models. ` +
|
||||
`Please check if the correct model is loaded/specified. ` +
|
||||
`Currently loaded models are:\n${loadedModels}`,
|
||||
);
|
||||
this.name = "SpecifiedModelNotFoundError";
|
||||
}
|
||||
}
|
||||
|
||||
export class IncorrectPipelineLoadedError extends Error {
|
||||
constructor(
|
||||
selectedModelId: string,
|
||||
expectedPipeline: string,
|
||||
requestName: string,
|
||||
) {
|
||||
super(
|
||||
`${requestName} expects model to be loaded with ${expectedPipeline}. However, ` +
|
||||
`${selectedModelId} is not loaded with this pipeline.`,
|
||||
);
|
||||
this.name = "IncorrectPipelineLoadedError";
|
||||
}
|
||||
}
|
||||
|
||||
export class ReloadModelIdNotUniqueError extends Error {
|
||||
constructor(modelId: string[]) {
|
||||
super(
|
||||
`Need to make models in modelId passed to reload() need to be unique. If you want to, ` +
|
||||
`load copies of the same model, consider making copies of the ModelRecord with ` +
|
||||
`different model_id. Received modelId: ${modelId}`,
|
||||
);
|
||||
this.name = "ReloadModelIdNotUniqueError";
|
||||
}
|
||||
}
|
||||
|
||||
export class IntegrityError extends Error {
|
||||
constructor(
|
||||
readonly url: string,
|
||||
readonly expected: string,
|
||||
readonly actual: string,
|
||||
) {
|
||||
super(
|
||||
`Integrity verification failed for ${url}\n` +
|
||||
` Expected: ${expected}\n` +
|
||||
` Actual: ${actual}\n` +
|
||||
`This may indicate file corruption or tampering.`,
|
||||
);
|
||||
this.name = "IntegrityError";
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
import * as tvmjs from "@mlc-ai/web-runtime";
|
||||
import log from "loglevel";
|
||||
import { ChatOptions, MLCEngineConfig } from "./config";
|
||||
import { ReloadParams, WorkerRequest } from "./message";
|
||||
import {
|
||||
ChatWorker,
|
||||
WebWorkerMLCEngineHandler,
|
||||
WebWorkerMLCEngine,
|
||||
} from "./web_worker";
|
||||
import { areArraysEqual, areChatOptionsListEqual } from "./utils";
|
||||
import { WebGPUNotFoundError } from "./error";
|
||||
|
||||
export interface ExtensionMLCEngineConfig extends MLCEngineConfig {
|
||||
extensionId?: string;
|
||||
onDisconnect?: () => void;
|
||||
}
|
||||
|
||||
/**
|
||||
* Worker handler that can be used in a ServiceWorker.
|
||||
*
|
||||
* @example
|
||||
*
|
||||
* const engine = new MLCEngine();
|
||||
* let handler;
|
||||
* chrome.runtime.onConnect.addListener(function (port) {
|
||||
* if (handler === undefined) {
|
||||
* handler = new ServiceWorkerMLCEngineHandler(engine, port);
|
||||
* } else {
|
||||
* handler.setPort(port);
|
||||
* }
|
||||
* port.onMessage.addListener(handler.onmessage.bind(handler));
|
||||
* });
|
||||
*/
|
||||
export class ServiceWorkerMLCEngineHandler extends WebWorkerMLCEngineHandler {
|
||||
port: chrome.runtime.Port | null;
|
||||
|
||||
constructor(port: chrome.runtime.Port) {
|
||||
super();
|
||||
this.port = port;
|
||||
port.onDisconnect.addListener(() => this.onPortDisconnect(port));
|
||||
}
|
||||
|
||||
postMessage(msg: any) {
|
||||
this.port?.postMessage(msg);
|
||||
}
|
||||
|
||||
setPort(port: chrome.runtime.Port) {
|
||||
this.port = port;
|
||||
port.onDisconnect.addListener(() => this.onPortDisconnect(port));
|
||||
}
|
||||
|
||||
onPortDisconnect(port: chrome.runtime.Port) {
|
||||
if (port === this.port) {
|
||||
this.port = null;
|
||||
}
|
||||
}
|
||||
|
||||
onmessage(event: any): void {
|
||||
if (event.type === "keepAlive") {
|
||||
return;
|
||||
}
|
||||
|
||||
const msg = event as WorkerRequest;
|
||||
if (msg.kind === "reload") {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as ReloadParams;
|
||||
// If the modelId, chatOpts, and appConfig are the same, immediately return
|
||||
if (
|
||||
areArraysEqual(this.modelId, params.modelId) &&
|
||||
areChatOptionsListEqual(this.chatOpts, params.chatOpts)
|
||||
) {
|
||||
log.info("Already loaded the model. Skip loading");
|
||||
const gpuDetectOutput = await tvmjs.detectGPUDevice();
|
||||
if (gpuDetectOutput == undefined) {
|
||||
throw new WebGPUNotFoundError();
|
||||
}
|
||||
let gpuLabel = "WebGPU";
|
||||
if (gpuDetectOutput.adapterInfo.description.length != 0) {
|
||||
gpuLabel += " - " + gpuDetectOutput.adapterInfo.description;
|
||||
} else {
|
||||
gpuLabel += " - " + gpuDetectOutput.adapterInfo.vendor;
|
||||
}
|
||||
this.engine.getInitProgressCallback()?.({
|
||||
progress: 1,
|
||||
timeElapsed: 0,
|
||||
text: "Finish loading on " + gpuLabel,
|
||||
});
|
||||
return null;
|
||||
}
|
||||
|
||||
await this.engine.reload(params.modelId, params.chatOpts);
|
||||
this.modelId = params.modelId;
|
||||
this.chatOpts = params.chatOpts;
|
||||
return null;
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
// All rest of message handling are the same as WebWorkerMLCEngineHandler
|
||||
super.onmessage(event);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a ServiceWorkerMLCEngine.
|
||||
*
|
||||
* @param modelId model_id of the model to load, either string or string[]. When multiple models
|
||||
* are provided, we load all models sequentially. Each modelId needs to either be in
|
||||
* `webllm.prebuiltAppConfig`, or in `engineCOnfig.appConfig`.
|
||||
* @param engineConfig Optionally configures the engine, see `webllm.MLCEngineConfig` for more.
|
||||
* @param chatOpts Extra options to optionally override the `mlc-chat-config.json` of `modelId`.
|
||||
* The size of which needs to match that of `modelId`; chatOpts[i] will be used for modelId[i].
|
||||
* @param keepAliveMs The interval to send keep alive messages to the service worker.
|
||||
* See [Service worker lifecycle](https://developer.chrome.com/docs/extensions/develop/concepts/service-workers/lifecycle#idle-shutdown)
|
||||
* The default is 10s.
|
||||
* @returns An initialized `WebLLM.ServiceWorkerMLCEngine` with `modelId` loaded.
|
||||
*/
|
||||
export async function CreateServiceWorkerMLCEngine(
|
||||
modelId: string | string[],
|
||||
engineConfig?: ExtensionMLCEngineConfig,
|
||||
chatOpts?: ChatOptions | ChatOptions[],
|
||||
keepAliveMs = 10000,
|
||||
): Promise<ServiceWorkerMLCEngine> {
|
||||
const serviceWorkerMLCEngine = new ServiceWorkerMLCEngine(
|
||||
engineConfig,
|
||||
keepAliveMs,
|
||||
);
|
||||
await serviceWorkerMLCEngine.reload(modelId, chatOpts);
|
||||
return serviceWorkerMLCEngine;
|
||||
}
|
||||
|
||||
class PortAdapter implements ChatWorker {
|
||||
port: chrome.runtime.Port;
|
||||
private _onmessage!: (message: any) => void;
|
||||
|
||||
constructor(port: chrome.runtime.Port) {
|
||||
this.port = port;
|
||||
this.port.onMessage.addListener(this.handleMessage.bind(this));
|
||||
}
|
||||
|
||||
// Wrapper to handle incoming messages and delegate to onmessage if available
|
||||
private handleMessage(message: any) {
|
||||
if (this._onmessage) {
|
||||
this._onmessage(message);
|
||||
}
|
||||
}
|
||||
|
||||
// Getter and setter for onmessage to manage adding/removing listeners
|
||||
get onmessage(): (message: any) => void {
|
||||
return this._onmessage;
|
||||
}
|
||||
|
||||
set onmessage(listener: (message: any) => void) {
|
||||
this._onmessage = listener;
|
||||
}
|
||||
|
||||
// Wrap port.postMessage to maintain 'this' context
|
||||
postMessage = (message: any): void => {
|
||||
this.port.postMessage(message);
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* A client of MLCEngine that exposes the same interface
|
||||
*/
|
||||
export class ServiceWorkerMLCEngine extends WebWorkerMLCEngine {
|
||||
port: chrome.runtime.Port;
|
||||
extensionId?: string;
|
||||
|
||||
constructor(engineConfig?: ExtensionMLCEngineConfig, keepAliveMs = 10000) {
|
||||
const extensionId = engineConfig?.extensionId;
|
||||
const onDisconnect = engineConfig?.onDisconnect;
|
||||
const port = extensionId
|
||||
? chrome.runtime.connect(extensionId, {
|
||||
name: "web_llm_service_worker",
|
||||
})
|
||||
: chrome.runtime.connect({ name: "web_llm_service_worker" });
|
||||
const chatWorker = new PortAdapter(port);
|
||||
super(chatWorker, engineConfig);
|
||||
this.port = port;
|
||||
this.extensionId = extensionId;
|
||||
|
||||
// Keep alive through periodical heartbeat signals
|
||||
const keepAliveTimer = setInterval(() => {
|
||||
this.worker.postMessage({ kind: "keepAlive" });
|
||||
}, keepAliveMs);
|
||||
|
||||
port.onDisconnect.addListener(() => {
|
||||
clearInterval(keepAliveTimer);
|
||||
if (onDisconnect) {
|
||||
onDisconnect();
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
export {
|
||||
ModelRecord,
|
||||
AppConfig,
|
||||
OPFSAccessMode,
|
||||
ChatOptions,
|
||||
MLCEngineConfig,
|
||||
GenerationConfig,
|
||||
ModelType,
|
||||
prebuiltAppConfig,
|
||||
modelVersion,
|
||||
modelLibURLPrefix,
|
||||
functionCallingModelIds,
|
||||
} from "./config";
|
||||
|
||||
export {
|
||||
verifyIntegrity,
|
||||
isValidSRI,
|
||||
type ModelIntegrity,
|
||||
type SRIString,
|
||||
type FileIntegrityMap,
|
||||
} from "./integrity";
|
||||
|
||||
export { IntegrityError } from "./error";
|
||||
|
||||
export {
|
||||
InitProgressCallback,
|
||||
InitProgressReport,
|
||||
MLCEngineInterface,
|
||||
LogitProcessor,
|
||||
LogLevel,
|
||||
} from "./types";
|
||||
|
||||
export { MLCEngine, CreateMLCEngine } from "./engine";
|
||||
|
||||
export {
|
||||
hasModelInCache,
|
||||
deleteChatConfigInCache,
|
||||
deleteModelAllInfoInCache,
|
||||
deleteModelWasmInCache,
|
||||
deleteModelInCache,
|
||||
} from "./cache_util";
|
||||
|
||||
export {
|
||||
WebWorkerMLCEngineHandler,
|
||||
WebWorkerMLCEngine,
|
||||
CreateWebWorkerMLCEngine,
|
||||
} from "./web_worker";
|
||||
|
||||
export { WorkerRequest, WorkerResponse, CustomRequestParams } from "./message";
|
||||
|
||||
export {
|
||||
ServiceWorkerMLCEngineHandler,
|
||||
ServiceWorkerMLCEngine,
|
||||
CreateServiceWorkerMLCEngine,
|
||||
} from "./service_worker";
|
||||
|
||||
export {
|
||||
ServiceWorkerMLCEngineHandler as ExtensionServiceWorkerMLCEngineHandler,
|
||||
ServiceWorkerMLCEngine as ExtensionServiceWorkerMLCEngine,
|
||||
CreateServiceWorkerMLCEngine as CreateExtensionServiceWorkerMLCEngine,
|
||||
} from "./extension_service_worker";
|
||||
|
||||
export * from "./openai_api_protocols/index";
|
||||
@@ -0,0 +1,150 @@
|
||||
import log from "loglevel";
|
||||
import { IntegrityError } from "./error";
|
||||
|
||||
/**
|
||||
* SRI (Subresource Integrity) hash string.
|
||||
* Format: "sha256-BASE64", "sha384-BASE64", or "sha512-BASE64".
|
||||
* See https://developer.mozilla.org/en-US/docs/Web/Security/Subresource_Integrity
|
||||
*/
|
||||
export type SRIString = string;
|
||||
|
||||
/** Map of filename to SRI hash for per-file verification. */
|
||||
export type FileIntegrityMap = Record<string, SRIString>;
|
||||
|
||||
/**
|
||||
* Integrity configuration for a model's artifacts.
|
||||
* All fields are optional — only specified artifacts will be verified.
|
||||
*
|
||||
* @param config SRI hash for the model's `mlc-chat-config.json`.
|
||||
* @param model_lib SRI hash for the WASM model library file.
|
||||
* @param tokenizer SRI hashes for tokenizer files, keyed by filename
|
||||
* (e.g. `"tokenizer.json"` or `"tokenizer.model"`).
|
||||
* @param onFailure Behavior on verification failure:
|
||||
* `"error"` (default) throws an `IntegrityError`;
|
||||
* `"warn"` logs a warning and continues.
|
||||
*/
|
||||
export interface ModelIntegrity {
|
||||
config?: SRIString;
|
||||
model_lib?: SRIString;
|
||||
tokenizer?: FileIntegrityMap;
|
||||
onFailure?: "error" | "warn";
|
||||
}
|
||||
|
||||
type SRIAlgorithm = "sha256" | "sha384" | "sha512";
|
||||
|
||||
const SRI_REGEX = /^(sha256|sha384|sha512)-([A-Za-z0-9+/]+={0,2})$/;
|
||||
|
||||
const SRI_HASH_BYTE_LENGTH: Record<SRIAlgorithm, number> = {
|
||||
sha256: 32,
|
||||
sha384: 48,
|
||||
sha512: 64,
|
||||
};
|
||||
|
||||
const ALGO_MAP: Record<SRIAlgorithm, string> = {
|
||||
sha256: "SHA-256",
|
||||
sha384: "SHA-384",
|
||||
sha512: "SHA-512",
|
||||
};
|
||||
|
||||
function getDecodedBase64ByteLength(base64: string): number | null {
|
||||
const length = base64.length;
|
||||
const remainder = length % 4;
|
||||
if (remainder === 1) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const paddingMatch = base64.match(/=+$/);
|
||||
const padding = paddingMatch ? paddingMatch[0].length : 0;
|
||||
if (padding > 2) {
|
||||
return null;
|
||||
}
|
||||
if (padding > 0 && remainder !== 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const fullQuartets = Math.floor(length / 4);
|
||||
let decodedLength = fullQuartets * 3;
|
||||
|
||||
if (padding > 0) {
|
||||
decodedLength -= padding;
|
||||
} else if (remainder === 2) {
|
||||
decodedLength += 1;
|
||||
} else if (remainder === 3) {
|
||||
decodedLength += 2;
|
||||
}
|
||||
|
||||
return decodedLength;
|
||||
}
|
||||
|
||||
function parseSRI(sri: string): { algo: SRIAlgorithm; hash: string } | null {
|
||||
const match = sri.match(SRI_REGEX);
|
||||
if (!match) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const algo = match[1] as SRIAlgorithm;
|
||||
const hash = match[2];
|
||||
const decodedLength = getDecodedBase64ByteLength(hash);
|
||||
if (decodedLength !== SRI_HASH_BYTE_LENGTH[algo]) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return { algo, hash };
|
||||
}
|
||||
|
||||
/**
|
||||
* Verify an ArrayBuffer against an SRI hash using the Web Crypto API.
|
||||
*
|
||||
* @param data The raw bytes to verify.
|
||||
* @param expectedSRI The expected SRI hash (e.g. `"sha256-abc123..."`).
|
||||
* @param url The URL of the artifact, used for error messages.
|
||||
* @param onFailure `"error"` to throw on mismatch, `"warn"` to log and continue.
|
||||
* @throws {IntegrityError} When the hash does not match and `onFailure` is `"error"`.
|
||||
*/
|
||||
export async function verifyIntegrity(
|
||||
data: ArrayBuffer,
|
||||
expectedSRI: SRIString,
|
||||
url: string,
|
||||
onFailure: "error" | "warn" = "error",
|
||||
): Promise<void> {
|
||||
const parsed = parseSRI(expectedSRI);
|
||||
if (!parsed) {
|
||||
throw new Error(
|
||||
`Invalid SRI hash format: "${expectedSRI}". ` +
|
||||
`Expected format: "sha256-BASE64", "sha384-BASE64", or "sha512-BASE64".`,
|
||||
);
|
||||
}
|
||||
|
||||
const { algo, hash: expectedHash } = parsed;
|
||||
const hashBuffer = await crypto.subtle.digest(ALGO_MAP[algo], data);
|
||||
const hashArray = new Uint8Array(hashBuffer);
|
||||
|
||||
// Convert to base64
|
||||
let binary = "";
|
||||
for (let i = 0; i < hashArray.length; i++) {
|
||||
binary += String.fromCharCode(hashArray[i]);
|
||||
}
|
||||
const actualHash = btoa(binary);
|
||||
|
||||
if (actualHash !== expectedHash) {
|
||||
const actualSRI = `${algo}-${actualHash}`;
|
||||
if (onFailure === "warn") {
|
||||
log.warn(
|
||||
`Integrity check failed for ${url}. ` +
|
||||
`Expected: ${expectedSRI}, Got: ${actualSRI}`,
|
||||
);
|
||||
return;
|
||||
}
|
||||
throw new IntegrityError(url, expectedSRI, actualSRI);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate that a string is a well-formed SRI hash.
|
||||
*
|
||||
* @param sri The string to validate.
|
||||
* @returns `true` if `sri` matches the SRI format.
|
||||
*/
|
||||
export function isValidSRI(sri: string): boolean {
|
||||
return parseSRI(sri) !== null;
|
||||
}
|
||||
+2286
File diff suppressed because it is too large
Load Diff
+163
@@ -0,0 +1,163 @@
|
||||
import { AppConfig, ChatOptions } from "./config";
|
||||
import { InitProgressReport, LogLevel } from "./types";
|
||||
import {
|
||||
ChatCompletionRequestStreaming,
|
||||
ChatCompletionRequestNonStreaming,
|
||||
ChatCompletion,
|
||||
ChatCompletionChunk,
|
||||
CompletionCreateParamsNonStreaming,
|
||||
CompletionCreateParamsStreaming,
|
||||
Completion,
|
||||
EmbeddingCreateParams,
|
||||
CreateEmbeddingResponse,
|
||||
} from "./openai_api_protocols/index";
|
||||
|
||||
/**
|
||||
* Message kind used by worker
|
||||
*/
|
||||
type RequestKind =
|
||||
| "reload"
|
||||
| "runtimeStatsText"
|
||||
| "interruptGenerate"
|
||||
| "unload"
|
||||
| "resetChat"
|
||||
| "getMaxStorageBufferBindingSize"
|
||||
| "getGPUVendor"
|
||||
| "forwardTokensAndSample"
|
||||
| "chatCompletionNonStreaming"
|
||||
| "completionNonStreaming"
|
||||
| "embedding"
|
||||
| "getMessage"
|
||||
| "chatCompletionStreamInit"
|
||||
| "completionStreamInit"
|
||||
| "completionStreamNextChunk"
|
||||
| "customRequest"
|
||||
| "keepAlive"
|
||||
| "setLogLevel"
|
||||
| "setAppConfig";
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
type ResponseKind = "return" | "throw" | "initProgressCallback";
|
||||
|
||||
export interface ReloadParams {
|
||||
modelId: string[];
|
||||
chatOpts?: ChatOptions[];
|
||||
}
|
||||
export interface ResetChatParams {
|
||||
keepStats: boolean;
|
||||
modelId?: string;
|
||||
}
|
||||
export interface GetMessageParams {
|
||||
modelId?: string;
|
||||
}
|
||||
export interface RuntimeStatsTextParams {
|
||||
modelId?: string;
|
||||
}
|
||||
export interface ForwardTokensAndSampleParams {
|
||||
inputIds: Array<number>;
|
||||
isPrefill: boolean;
|
||||
modelId?: string;
|
||||
}
|
||||
|
||||
// Notes on the following Params with modelId and chatOpts:
|
||||
// These fields are the model and chatOpts that the frontend engine expects the backend
|
||||
// to be loaded with. If not loaded due to web/service worker unexpectedly killed,
|
||||
// handler will call reload(). An engine can load multiple models, hence both are list.
|
||||
// TODO(webllm-team): should add appConfig here as well if rigorous.
|
||||
// Fore more, see https://github.com/mlc-ai/web-llm/pull/471
|
||||
|
||||
// Note on the messages with selectedModelId:
|
||||
// This is the modelId this request uses. It is needed to identify which async generator
|
||||
// to instantiate / use, since an engine can load multiple models, thus the handler
|
||||
// needs to maintain multiple generators.
|
||||
export interface ChatCompletionNonStreamingParams {
|
||||
request: ChatCompletionRequestNonStreaming;
|
||||
modelId: string[];
|
||||
chatOpts?: ChatOptions[];
|
||||
}
|
||||
export interface ChatCompletionStreamInitParams {
|
||||
request: ChatCompletionRequestStreaming;
|
||||
selectedModelId: string;
|
||||
modelId: string[];
|
||||
chatOpts?: ChatOptions[];
|
||||
}
|
||||
export interface CompletionNonStreamingParams {
|
||||
request: CompletionCreateParamsNonStreaming;
|
||||
modelId: string[];
|
||||
chatOpts?: ChatOptions[];
|
||||
}
|
||||
export interface CompletionStreamInitParams {
|
||||
request: CompletionCreateParamsStreaming;
|
||||
selectedModelId: string;
|
||||
modelId: string[];
|
||||
chatOpts?: ChatOptions[];
|
||||
}
|
||||
export interface EmbeddingParams {
|
||||
request: EmbeddingCreateParams;
|
||||
modelId: string[];
|
||||
chatOpts?: ChatOptions[];
|
||||
}
|
||||
export interface CompletionStreamNextChunkParams {
|
||||
selectedModelId: string;
|
||||
}
|
||||
|
||||
export interface CustomRequestParams {
|
||||
requestName: string;
|
||||
requestMessage: string;
|
||||
}
|
||||
export type MessageContent =
|
||||
| ReloadParams
|
||||
| ResetChatParams
|
||||
| GetMessageParams
|
||||
| RuntimeStatsTextParams
|
||||
| ForwardTokensAndSampleParams
|
||||
| ChatCompletionNonStreamingParams
|
||||
| ChatCompletionStreamInitParams
|
||||
| CompletionNonStreamingParams
|
||||
| CompletionStreamInitParams
|
||||
| EmbeddingParams
|
||||
| CompletionStreamNextChunkParams
|
||||
| CustomRequestParams
|
||||
| InitProgressReport
|
||||
| LogLevel
|
||||
| string
|
||||
| null
|
||||
| number
|
||||
| ChatCompletion
|
||||
| ChatCompletionChunk
|
||||
| CreateEmbeddingResponse
|
||||
| Completion
|
||||
| AppConfig
|
||||
| void;
|
||||
/**
|
||||
* The message used in exchange between worker
|
||||
* and the main thread.
|
||||
*/
|
||||
|
||||
export type WorkerRequest = {
|
||||
kind: RequestKind;
|
||||
uuid: string;
|
||||
content: MessageContent;
|
||||
};
|
||||
|
||||
type HeartbeatWorkerResponse = {
|
||||
kind: "heartbeat";
|
||||
uuid: string;
|
||||
};
|
||||
|
||||
type OneTimeWorkerResponse = {
|
||||
kind: "return" | "throw";
|
||||
uuid: string;
|
||||
content: MessageContent;
|
||||
};
|
||||
|
||||
type InitProgressWorkerResponse = {
|
||||
kind: "initProgressCallback";
|
||||
uuid: string;
|
||||
content: InitProgressReport;
|
||||
};
|
||||
|
||||
export type WorkerResponse =
|
||||
| OneTimeWorkerResponse
|
||||
| InitProgressWorkerResponse
|
||||
| HeartbeatWorkerResponse;
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,381 @@
|
||||
/**
|
||||
* The input to OpenAI API, directly adopted from openai-node with small tweaks:
|
||||
* https://github.com/openai/openai-node/blob/master/src/resources/completions.ts
|
||||
*
|
||||
* Copyright 2024 OpenAI
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
import { MLCEngineInterface } from "../types";
|
||||
import {
|
||||
InvalidStreamOptionsError,
|
||||
SeedTypeError,
|
||||
StreamingCountError,
|
||||
UnsupportedFieldsError,
|
||||
} from "../error";
|
||||
import {
|
||||
ChatCompletion,
|
||||
ChatCompletionStreamOptions,
|
||||
CompletionUsage,
|
||||
ChatCompletionFinishReason,
|
||||
} from "./chat_completion";
|
||||
|
||||
export class Completions {
|
||||
private engine: MLCEngineInterface;
|
||||
|
||||
constructor(engine: MLCEngineInterface) {
|
||||
this.engine = engine;
|
||||
}
|
||||
|
||||
create(request: CompletionCreateParamsNonStreaming): Promise<Completion>;
|
||||
create(
|
||||
request: CompletionCreateParamsStreaming,
|
||||
): Promise<AsyncIterable<Completion>>;
|
||||
create(
|
||||
request: CompletionCreateParamsBase,
|
||||
): Promise<AsyncIterable<Completion> | Completion>;
|
||||
create(
|
||||
request: CompletionCreateParams,
|
||||
): Promise<AsyncIterable<Completion> | Completion> {
|
||||
return this.engine.completion(request);
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////// 1. CREATE PARAMS ////////////////////////////////
|
||||
/**
|
||||
* OpenAI completion request protocol.
|
||||
*
|
||||
* API reference: https://platform.openai.com/docs/api-reference/completions/create
|
||||
* Followed: https://github.com/openai/openai-node/blob/master/src/resources/completions.ts
|
||||
*
|
||||
* @note `model` is excluded. Instead, call `CreateMLCEngine(model)` or `engine.reload(model)` explicitly before calling this API.
|
||||
*/
|
||||
export interface CompletionCreateParamsBase {
|
||||
/**
|
||||
* The prompt(s) to generate completions for, encoded as a string.
|
||||
*/
|
||||
prompt: string;
|
||||
|
||||
/**
|
||||
* Echo back the prompt in addition to the completion
|
||||
*/
|
||||
echo?: boolean | null;
|
||||
|
||||
/**
|
||||
* Number between -2.0 and 2.0. Positive values penalize new tokens based on their
|
||||
* existing frequency in the text so far, decreasing the model's likelihood to
|
||||
* repeat the same line verbatim.
|
||||
*
|
||||
* [See more information about frequency and presence penalties.](https://platform.openai.com/docs/guides/text-generation/parameter-details)
|
||||
*/
|
||||
frequency_penalty?: number | null;
|
||||
|
||||
/**
|
||||
* Modify the likelihood of specified tokens appearing in the completion.
|
||||
*
|
||||
* Accepts a JSON object that maps tokens (specified by their token ID, which varies per model)
|
||||
* to an associated bias value from -100 to 100. Typically, you can see `tokenizer.json` of the
|
||||
* model to see which token ID maps to what string. Mathematically, the bias is added to the
|
||||
* logits generated by the model prior to sampling. The exact effect will vary per model, but
|
||||
* values between -1 and 1 should decrease or increase likelihood of selection; values like -100
|
||||
* or 100 should result in a ban or exclusive selection of the relevant token.
|
||||
*
|
||||
* As an example, you can pass `{"16230": -100}` to prevent the `Hello` token from being
|
||||
* generated in Mistral-7B-Instruct-v0.2, according to the mapping in
|
||||
* https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2/raw/main/tokenizer.json.
|
||||
*
|
||||
* @note For stateful and customizable / flexible logit processing, see `webllm.LogitProcessor`.
|
||||
* @note If used in combination with `webllm.LogitProcessor`, `logit_bias` is applied after
|
||||
* `LogitProcessor.processLogits()` is called.
|
||||
*/
|
||||
logit_bias?: Record<string, number> | null;
|
||||
|
||||
/**
|
||||
* Whether to return log probabilities of the output tokens or not.
|
||||
*
|
||||
* If true, returns the log probabilities of each output token returned in the `content` of
|
||||
* `message`.
|
||||
*/
|
||||
logprobs?: boolean | null;
|
||||
|
||||
/**
|
||||
* An integer between 0 and 5 specifying the number of most likely tokens to return
|
||||
* at each token position, each with an associated log probability. `logprobs` must
|
||||
* be set to `true` if this parameter is used.
|
||||
*/
|
||||
top_logprobs?: number | null;
|
||||
|
||||
/**
|
||||
* The maximum number of [tokens](/tokenizer) that can be generated in the
|
||||
* completion.
|
||||
*
|
||||
* The total length of input tokens and generated tokens is limited by the model's
|
||||
* context length.
|
||||
*/
|
||||
max_tokens?: number | null;
|
||||
|
||||
/**
|
||||
* How many completions to generate for each prompt.
|
||||
*/
|
||||
n?: number | null;
|
||||
|
||||
/**
|
||||
* Number between -2.0 and 2.0. Positive values penalize new tokens based on
|
||||
* whether they appear in the text so far, increasing the model's likelihood to
|
||||
* talk about new topics.
|
||||
*
|
||||
* [See more information about frequency and presence penalties.](https://platform.openai.com/docs/guides/text-generation/parameter-details)
|
||||
*/
|
||||
presence_penalty?: number | null;
|
||||
|
||||
/**
|
||||
* Penalizes new tokens based on whether they appear in the prompt and the
|
||||
* generated text so far. Values greater than 1.0 encourage the model to use new
|
||||
* tokens, while values less than 1.0 encourage the model to repeat tokens.
|
||||
*/
|
||||
repetition_penalty?: number | null;
|
||||
|
||||
/**
|
||||
* If specified, our system will make a best effort to sample deterministically,
|
||||
* such that repeated requests with the same `seed` and parameters should return
|
||||
* the same result.
|
||||
*
|
||||
* @note Seeding is done on a request-level rather than choice-level. That is, if `n > 1`, you
|
||||
* would still get different content for each `Choice`. But if two requests with `n = 2` are
|
||||
* processed with the same seed, the two results should be the same (two choices are different).
|
||||
*/
|
||||
seed?: number | null;
|
||||
|
||||
/**
|
||||
* Up to 4 sequences where the API will stop generating further tokens. The
|
||||
* returned text will not contain the stop sequence.
|
||||
*/
|
||||
stop?: string | null | Array<string>;
|
||||
|
||||
/**
|
||||
* If set, partial deltas will be sent. It will be terminated by an empty chunk.
|
||||
*/
|
||||
stream?: boolean | null;
|
||||
|
||||
/**
|
||||
* Options for streaming response. Only set this when you set `stream: true`.
|
||||
*/
|
||||
stream_options?: ChatCompletionStreamOptions | null;
|
||||
|
||||
/**
|
||||
* What sampling temperature to use, between 0 and 2. Higher values like 0.8 will
|
||||
* make the output more random, while lower values like 0.2 will make it more
|
||||
* focused and deterministic.
|
||||
*
|
||||
* We generally recommend altering this or `top_p` but not both.
|
||||
*/
|
||||
temperature?: number | null;
|
||||
|
||||
/**
|
||||
* An alternative to sampling with temperature, called nucleus sampling, where the
|
||||
* model considers the results of the tokens with top_p probability mass. So 0.1
|
||||
* means only the tokens comprising the top 10% probability mass are considered.
|
||||
*
|
||||
* We generally recommend altering this or `temperature` but not both.
|
||||
*/
|
||||
top_p?: number | null;
|
||||
|
||||
/**
|
||||
* If true, will ignore stop string and stop token and generate until max_tokens hit.
|
||||
* If unset, will treat as false.
|
||||
*/
|
||||
ignore_eos?: boolean;
|
||||
|
||||
/**
|
||||
* ID of the model to use. This equals to `ModelRecord.model_id`, which needs to either be in
|
||||
* `webllm.prebuiltAppConfig` or in `engineConfig.appConfig`.
|
||||
*
|
||||
* @note Call `CreateMLCEngine(model)` or `engine.reload(model)` ahead of time.
|
||||
* @note If only one model is loaded in the engine, this field is optional. If multiple models
|
||||
* are loaded, this is required.
|
||||
*/
|
||||
model?: string | null;
|
||||
|
||||
//////////////// BELOW FIELDS NOT SUPPORTED YET ////////////////
|
||||
|
||||
/**
|
||||
* The suffix that comes after a completion of inserted text.
|
||||
*
|
||||
* @note This field is not supported.
|
||||
*/
|
||||
suffix?: string | null;
|
||||
|
||||
/**
|
||||
* A unique identifier representing your end-user, which can help OpenAI to monitor
|
||||
* and detect abuse.
|
||||
*
|
||||
* @note This field is not supported.
|
||||
*/
|
||||
user?: string;
|
||||
|
||||
/**
|
||||
* Generates `best_of` completions server-side and returns the "best" (the one with
|
||||
* the highest log probability per token). Results cannot be streamed.
|
||||
*
|
||||
* When used with `n`, `best_of` controls the number of candidate completions and
|
||||
* `n` specifies how many to return – `best_of` must be greater than `n`.
|
||||
*
|
||||
* @note This field is not supported.
|
||||
*/
|
||||
best_of?: number | null;
|
||||
|
||||
/**
|
||||
* Fields specific to WebLLM, not present in OpenAI.
|
||||
*/
|
||||
extra_body?: {
|
||||
/**
|
||||
* If set to true, the response will include a breakdown of the time spent in various
|
||||
* stages of token sampling.
|
||||
*/
|
||||
enable_latency_breakdown?: boolean | null;
|
||||
};
|
||||
}
|
||||
|
||||
export type CompletionCreateParams =
|
||||
| CompletionCreateParamsNonStreaming
|
||||
| CompletionCreateParamsStreaming;
|
||||
|
||||
export interface CompletionCreateParamsNonStreaming
|
||||
extends CompletionCreateParamsBase {
|
||||
/**
|
||||
* If set, partial deltas will be sent. It will be terminated by an empty chunk.
|
||||
*/
|
||||
stream?: false | null;
|
||||
}
|
||||
|
||||
export interface CompletionCreateParamsStreaming
|
||||
extends CompletionCreateParamsBase {
|
||||
/**
|
||||
* If set, partial deltas will be sent. It will be terminated by an empty chunk.
|
||||
*/
|
||||
stream: true;
|
||||
}
|
||||
|
||||
//////////////////////////////// 2. RESPONSE ////////////////////////////////
|
||||
/**
|
||||
* Represents a completion response returned by model, based on the provided input.
|
||||
*/
|
||||
export interface Completion {
|
||||
/**
|
||||
* A unique identifier for the completion.
|
||||
*/
|
||||
id: string;
|
||||
|
||||
/**
|
||||
* The list of completion choices the model generated for the input prompt.
|
||||
*/
|
||||
choices: Array<CompletionChoice>;
|
||||
|
||||
/**
|
||||
* The Unix timestamp (in seconds) of when the completion was created.
|
||||
*/
|
||||
created: number;
|
||||
|
||||
/**
|
||||
* The model used for completion.
|
||||
*/
|
||||
model: string;
|
||||
|
||||
/**
|
||||
* The object type, which is always "text_completion"
|
||||
*/
|
||||
object: "text_completion";
|
||||
|
||||
/**
|
||||
* This fingerprint represents the backend configuration that the model runs with.
|
||||
*
|
||||
* Can be used in conjunction with the `seed` request parameter to understand when
|
||||
* backend changes have been made that might impact determinism.
|
||||
*
|
||||
* @note Not supported yet.
|
||||
*/
|
||||
system_fingerprint?: string;
|
||||
|
||||
/**
|
||||
* Usage statistics for the completion request.
|
||||
*/
|
||||
usage?: CompletionUsage;
|
||||
}
|
||||
|
||||
export interface CompletionChoice {
|
||||
/**
|
||||
* The reason the model stopped generating tokens. This will be `stop` if the model
|
||||
* hit a natural stop point or a provided stop sequence, or `length` if the maximum
|
||||
* number of tokens specified in the request was reached.
|
||||
*/
|
||||
finish_reason: ChatCompletionFinishReason | null;
|
||||
|
||||
index: number;
|
||||
|
||||
/**
|
||||
* A list of message content tokens with log probability information.
|
||||
* @note Different from openai-node, we reuse ChatCompletion's Logprobs.
|
||||
*/
|
||||
logprobs?: ChatCompletion.Choice.Logprobs | null;
|
||||
|
||||
text: string;
|
||||
}
|
||||
|
||||
//////////////////////////////// 3. POST INIT ////////////////////////////////
|
||||
|
||||
export const CompletionCreateParamsUnsupportedFields: Array<string> = [
|
||||
"suffix",
|
||||
"user",
|
||||
"best_of",
|
||||
];
|
||||
|
||||
/**
|
||||
* Post init and verify whether the input of the request is valid. Thus, this function can throw
|
||||
* error or in-place update request.
|
||||
* @param request User's input request.
|
||||
* @param currentModelId The current model loaded that will perform this request.
|
||||
*/
|
||||
export function postInitAndCheckFields(
|
||||
request: CompletionCreateParams,
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
currentModelId: string,
|
||||
): void {
|
||||
// 1. Check unsupported fields in request
|
||||
const unsupported: Array<string> = [];
|
||||
CompletionCreateParamsUnsupportedFields.forEach((field) => {
|
||||
if (field in request) {
|
||||
unsupported.push(field);
|
||||
}
|
||||
});
|
||||
if (unsupported.length > 0) {
|
||||
throw new UnsupportedFieldsError(unsupported, "CompletionCreateParams");
|
||||
}
|
||||
|
||||
// 2. If streaming, n cannot be > 1, since we cannot manage multiple sequences at once
|
||||
if (request.stream && request.n && request.n > 1) {
|
||||
throw new StreamingCountError();
|
||||
}
|
||||
|
||||
// 3. Seed should be an integer
|
||||
if (request.seed !== undefined && request.seed !== null) {
|
||||
if (!Number.isInteger(request.seed)) {
|
||||
throw new SeedTypeError(request.seed);
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Only set stream_options when streaming
|
||||
if (request.stream_options !== undefined && request.stream_options !== null) {
|
||||
if (!request.stream) {
|
||||
throw new InvalidStreamOptionsError();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
/**
|
||||
* The input to OpenAI API, directly adopted from openai-node with small tweaks:
|
||||
* https://github.com/openai/openai-node/blob/master/src/resources/embeddings.ts
|
||||
*
|
||||
* Copyright 2024 OpenAI
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
import {
|
||||
EmbeddingInputEmptyError,
|
||||
EmbeddingUnsupportedEncodingFormatError,
|
||||
UnsupportedFieldsError,
|
||||
} from "../error";
|
||||
import { MLCEngineInterface } from "../types";
|
||||
|
||||
export class Embeddings {
|
||||
private engine: MLCEngineInterface;
|
||||
|
||||
constructor(engine: MLCEngineInterface) {
|
||||
this.engine = engine;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates an embedding vector representing the input text.
|
||||
*/
|
||||
create(request: EmbeddingCreateParams): Promise<CreateEmbeddingResponse> {
|
||||
return this.engine.embedding(request);
|
||||
}
|
||||
}
|
||||
|
||||
export interface CreateEmbeddingResponse {
|
||||
/**
|
||||
* The list of embeddings generated by the model.
|
||||
*/
|
||||
data: Array<Embedding>;
|
||||
|
||||
/**
|
||||
* The name of the model used to generate the embedding.
|
||||
*/
|
||||
model: string;
|
||||
|
||||
/**
|
||||
* The object type, which is always "list".
|
||||
*/
|
||||
object: "list";
|
||||
|
||||
/**
|
||||
* The usage information for the request.
|
||||
*/
|
||||
usage: CreateEmbeddingResponse.Usage;
|
||||
}
|
||||
|
||||
/* eslint-disable-next-line @typescript-eslint/no-namespace */
|
||||
export namespace CreateEmbeddingResponse {
|
||||
/**
|
||||
* The usage information for the request.
|
||||
*/
|
||||
export interface Usage {
|
||||
/**
|
||||
* The number of tokens used by the prompt.
|
||||
*/
|
||||
prompt_tokens: number;
|
||||
|
||||
/**
|
||||
* The total number of tokens used by the request.
|
||||
*/
|
||||
total_tokens: number;
|
||||
|
||||
/**
|
||||
* Fields specific to WebLLM, not present in OpenAI.
|
||||
*/
|
||||
extra: {
|
||||
/**
|
||||
* Number of tokens per second for prefilling.
|
||||
*/
|
||||
prefill_tokens_per_s: number;
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents an embedding vector returned by embedding endpoint.
|
||||
*/
|
||||
export interface Embedding {
|
||||
/**
|
||||
* The embedding vector, which is a list of floats. The length of vector depends on
|
||||
* the model.
|
||||
*/
|
||||
embedding: Array<number>;
|
||||
|
||||
/**
|
||||
* The index of the embedding in the list of embeddings.
|
||||
*/
|
||||
index: number;
|
||||
|
||||
/**
|
||||
* The object type, which is always "embedding".
|
||||
*/
|
||||
object: "embedding";
|
||||
}
|
||||
|
||||
export interface EmbeddingCreateParams {
|
||||
/**
|
||||
* Input text to embed, encoded as a string or array of tokens. To embed multiple
|
||||
* inputs in a single request, pass an array of strings or array of token arrays.
|
||||
* The input must not exceed the max input tokens for the model, and cannot be an empty string.
|
||||
* If the batch size is too large, multiple forward of the will take place.
|
||||
*/
|
||||
input: string | Array<string> | Array<number> | Array<Array<number>>;
|
||||
|
||||
/**
|
||||
* ID of the model to use. This equals to `ModelRecord.model_id`, which needs to either be in
|
||||
* `webllm.prebuiltAppConfig` or in `engineConfig.appConfig`.
|
||||
*
|
||||
* @note Call `CreateMLCEngine(model)` or `engine.reload(model)` ahead of time.
|
||||
* @note If only one model is loaded in the engine, this field is optional. If multiple models
|
||||
* are loaded, this is required.
|
||||
*/
|
||||
model?: string | null;
|
||||
|
||||
/**
|
||||
* The format to return the embeddings in.
|
||||
*
|
||||
* @note Currently only support `float`.
|
||||
*/
|
||||
encoding_format?: "float" | "base64";
|
||||
|
||||
// TODO: can support matryoshka embedding models in future, hence allow `dimensions` for those.
|
||||
/**
|
||||
* The number of dimensions the resulting output embeddings should have.
|
||||
*
|
||||
* @note Not supported.
|
||||
*/
|
||||
dimensions?: number;
|
||||
|
||||
/**
|
||||
* A unique identifier representing your end-user, which can help OpenAI to monitor
|
||||
* and detect abuse.
|
||||
*
|
||||
* @note Not supported.
|
||||
*/
|
||||
user?: string;
|
||||
}
|
||||
|
||||
export const EmbeddingCreateParamsUnsupportedFields: Array<string> = [
|
||||
"dimensions",
|
||||
"user",
|
||||
];
|
||||
|
||||
export function postInitAndCheckFields(
|
||||
request: EmbeddingCreateParams,
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
currentModelId: string,
|
||||
): void {
|
||||
// 1. Check unsupported fields in request
|
||||
const unsupported: Array<string> = [];
|
||||
EmbeddingCreateParamsUnsupportedFields.forEach((field) => {
|
||||
if (field in request) {
|
||||
unsupported.push(field);
|
||||
}
|
||||
});
|
||||
if (unsupported.length > 0) {
|
||||
throw new UnsupportedFieldsError(unsupported, "EmbeddingCreateParams");
|
||||
}
|
||||
|
||||
// 2. Unsupported format
|
||||
if (request.encoding_format == "base64") {
|
||||
throw new EmbeddingUnsupportedEncodingFormatError();
|
||||
}
|
||||
|
||||
// 3. Invalid input
|
||||
const input = request.input;
|
||||
if (typeof input === "string") {
|
||||
if (input === "") throw new EmbeddingInputEmptyError();
|
||||
} else {
|
||||
// input instanceof Array
|
||||
if (input.length === 0) {
|
||||
// Array<number>
|
||||
throw new EmbeddingInputEmptyError();
|
||||
}
|
||||
for (let i = 0; i < input.length; i++) {
|
||||
const curInput = input[i];
|
||||
if (typeof curInput !== "number") {
|
||||
// Array<string>, Array<Array<number>>
|
||||
if (curInput.length === 0) throw new EmbeddingInputEmptyError();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
/**
|
||||
* The input to OpenAI API, directly adopted from openai-node with small tweaks:
|
||||
* https://github.com/openai/openai-node/blob/master/src/resources/chat/completions.ts
|
||||
*
|
||||
* Copyright 2024 OpenAI
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
export {
|
||||
Chat,
|
||||
ChatCompletionRequestBase,
|
||||
ChatCompletionRequestNonStreaming,
|
||||
ChatCompletionRequestStreaming,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletion,
|
||||
ChatCompletionChunk,
|
||||
ChatCompletionRequestUnsupportedFields,
|
||||
postInitAndCheckFields as postInitAndCheckFieldsChatCompletion,
|
||||
ChatCompletionContentPart,
|
||||
ChatCompletionContentPartText,
|
||||
ChatCompletionContentPartImage,
|
||||
ChatCompletionMessageToolCall,
|
||||
ChatCompletionRole,
|
||||
ChatCompletionSystemMessageParam,
|
||||
ChatCompletionUserMessageParam,
|
||||
ChatCompletionAssistantMessageParam,
|
||||
ChatCompletionToolMessageParam,
|
||||
ChatCompletionMessageParam,
|
||||
FunctionParameters,
|
||||
FunctionDefinition,
|
||||
ChatCompletionTool,
|
||||
ChatCompletionNamedToolChoice,
|
||||
ChatCompletionToolChoiceOption,
|
||||
TopLogprob,
|
||||
ChatCompletionTokenLogprob,
|
||||
ChatCompletionMessage,
|
||||
CompletionUsage,
|
||||
ResponseFormat,
|
||||
ChatCompletionFinishReason,
|
||||
} from "./chat_completion";
|
||||
|
||||
export {
|
||||
Completions,
|
||||
CompletionCreateParamsNonStreaming,
|
||||
CompletionCreateParamsStreaming,
|
||||
CompletionCreateParamsBase,
|
||||
CompletionCreateParams,
|
||||
Completion,
|
||||
CompletionChoice,
|
||||
postInitAndCheckFields as postInitAndCheckFieldsCompletion,
|
||||
} from "./completion";
|
||||
|
||||
export {
|
||||
Embeddings,
|
||||
Embedding,
|
||||
EmbeddingCreateParams,
|
||||
CreateEmbeddingResponse,
|
||||
postInitAndCheckFields as postInitAndCheckFieldsEmbedding,
|
||||
} from "./embedding";
|
||||
@@ -0,0 +1,253 @@
|
||||
import * as tvmjs from "@mlc-ai/web-runtime";
|
||||
import log from "loglevel";
|
||||
import { ChatOptions, MLCEngineConfig } from "./config";
|
||||
import { ReloadParams, WorkerRequest, WorkerResponse } from "./message";
|
||||
import { InitProgressReport } from "./types";
|
||||
import {
|
||||
WebWorkerMLCEngineHandler,
|
||||
WebWorkerMLCEngine,
|
||||
ChatWorker,
|
||||
} from "./web_worker";
|
||||
import { areArraysEqual, areChatOptionsListEqual } from "./utils";
|
||||
import {
|
||||
NoServiceWorkerAPIError,
|
||||
NonWorkerEnvironmentError,
|
||||
ServiceWorkerInitializationError,
|
||||
} from "./error";
|
||||
|
||||
/* Service Worker Script */
|
||||
|
||||
type IServiceWorker = globalThis.ServiceWorker;
|
||||
|
||||
/**
|
||||
* Worker handler that can be used in a ServiceWorker.
|
||||
*
|
||||
* @example
|
||||
*
|
||||
* const engine = new MLCEngine();
|
||||
* let handler;
|
||||
* chrome.runtime.onConnect.addListener(function (port) {
|
||||
* if (handler === undefined) {
|
||||
* handler = new ServiceWorkerMLCEngineHandler(engine, port);
|
||||
* } else {
|
||||
* handler.setPort(port);
|
||||
* }
|
||||
* port.onMessage.addListener(handler.onmessage.bind(handler));
|
||||
* });
|
||||
*/
|
||||
export class ServiceWorkerMLCEngineHandler extends WebWorkerMLCEngineHandler {
|
||||
private clientRegistry = new Map<
|
||||
string,
|
||||
IServiceWorker | Client | MessagePort
|
||||
>();
|
||||
private initRequestUuid?: string;
|
||||
|
||||
constructor() {
|
||||
if (!self || !("addEventListener" in self)) {
|
||||
throw new NonWorkerEnvironmentError("ServiceWorkerMLCEngineHandler");
|
||||
}
|
||||
super();
|
||||
const onmessage = this.onmessage.bind(this);
|
||||
|
||||
this.engine.setInitProgressCallback((report: InitProgressReport) => {
|
||||
const msg: WorkerResponse = {
|
||||
kind: "initProgressCallback",
|
||||
uuid: this.initRequestUuid || "",
|
||||
content: report,
|
||||
};
|
||||
this.postMessage(msg);
|
||||
});
|
||||
|
||||
self.addEventListener("message", (event) => {
|
||||
const message = event as unknown as ExtendableMessageEvent;
|
||||
if (message.source) {
|
||||
this.clientRegistry.set(message.data.uuid, message.source);
|
||||
}
|
||||
message.waitUntil(
|
||||
new Promise((resolve, reject) => {
|
||||
onmessage(message, resolve, reject);
|
||||
}),
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
postMessage(message: WorkerResponse) {
|
||||
if (this.clientRegistry.has(message.uuid)) {
|
||||
const client = this.clientRegistry.get(message.uuid);
|
||||
client?.postMessage(message);
|
||||
|
||||
if (message.kind === "return" || message.kind === "throw") {
|
||||
this.clientRegistry.delete(message.uuid);
|
||||
} else {
|
||||
// TODO(nestor): Delete clientRegistry after complete to avoid memory leak?
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
onmessage(
|
||||
event: ExtendableMessageEvent,
|
||||
onComplete?: (value: any) => void,
|
||||
onError?: () => void,
|
||||
): void {
|
||||
const msg = event.data as WorkerRequest;
|
||||
log.trace(
|
||||
`ServiceWorker message: [${msg.kind}] ${JSON.stringify(msg.content)}`,
|
||||
);
|
||||
|
||||
// Special case message handling different from WebWorkerMLCEngineHandler
|
||||
if (msg.kind === "keepAlive") {
|
||||
const reply: WorkerResponse = {
|
||||
kind: "heartbeat",
|
||||
uuid: msg.uuid,
|
||||
};
|
||||
this.postMessage(reply);
|
||||
onComplete?.(reply);
|
||||
return;
|
||||
}
|
||||
|
||||
if (msg.kind === "reload") {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as ReloadParams;
|
||||
// If the modelId, chatOpts, and appConfig are the same, immediately return
|
||||
if (
|
||||
areArraysEqual(this.modelId, params.modelId) &&
|
||||
areChatOptionsListEqual(this.chatOpts, params.chatOpts)
|
||||
) {
|
||||
log.info("Already loaded the model. Skip loading");
|
||||
const gpuDetectOutput = await tvmjs.detectGPUDevice();
|
||||
if (gpuDetectOutput == undefined) {
|
||||
throw Error("Cannot find WebGPU in the environment");
|
||||
}
|
||||
let gpuLabel = "WebGPU";
|
||||
if (gpuDetectOutput.adapterInfo.description.length != 0) {
|
||||
gpuLabel += " - " + gpuDetectOutput.adapterInfo.description;
|
||||
} else {
|
||||
gpuLabel += " - " + gpuDetectOutput.adapterInfo.vendor;
|
||||
}
|
||||
this.engine.getInitProgressCallback()?.({
|
||||
progress: 1,
|
||||
timeElapsed: 0,
|
||||
text: "Finish loading on " + gpuLabel,
|
||||
});
|
||||
onComplete?.(null);
|
||||
return null;
|
||||
}
|
||||
|
||||
this.initRequestUuid = msg.uuid;
|
||||
await this.engine.reload(params.modelId, params.chatOpts);
|
||||
this.modelId = params.modelId;
|
||||
this.chatOpts = params.chatOpts;
|
||||
onComplete?.(null);
|
||||
return null;
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
// All rest of message handling are the same as WebWorkerMLCEngineHandler
|
||||
super.onmessage(msg, onComplete, onError);
|
||||
}
|
||||
}
|
||||
|
||||
/* Webapp Client */
|
||||
export class ServiceWorker implements ChatWorker {
|
||||
_onmessage: (event: MessageEvent) => void = () => {};
|
||||
|
||||
get onmessage() {
|
||||
return this._onmessage;
|
||||
}
|
||||
|
||||
set onmessage(handler: (event: any) => void) {
|
||||
this._onmessage = handler;
|
||||
|
||||
if (!("serviceWorker" in navigator)) {
|
||||
throw new NoServiceWorkerAPIError();
|
||||
}
|
||||
(navigator.serviceWorker as ServiceWorkerContainer).onmessage = handler;
|
||||
}
|
||||
|
||||
postMessage(message: WorkerRequest) {
|
||||
if (!("serviceWorker" in navigator)) {
|
||||
throw new NoServiceWorkerAPIError();
|
||||
}
|
||||
const serviceWorker = (navigator.serviceWorker as ServiceWorkerContainer)
|
||||
.controller;
|
||||
if (!serviceWorker) {
|
||||
throw new Error("There is no active service worker");
|
||||
}
|
||||
serviceWorker.postMessage(message);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a ServiceWorkerMLCEngine.
|
||||
*
|
||||
* @param modelId model_id of the model to load, either string or string[]. When multiple models
|
||||
* are provided, we load all models sequentially. Each modelId needs to either be in
|
||||
* `webllm.prebuiltAppConfig`, or in `engineCOnfig.appConfig`.
|
||||
* @param engineConfig Optionally configures the engine, see `webllm.MLCEngineConfig` for more.
|
||||
* @param chatOpts Extra options to optionally override the `mlc-chat-config.json` of `modelId`.
|
||||
* The size of which needs to match that of `modelId`; chatOpts[i] will be used for modelId[i].
|
||||
* @returns An initialized `WebLLM.ServiceWorkerMLCEngine` with `modelId` loaded.
|
||||
*/
|
||||
export async function CreateServiceWorkerMLCEngine(
|
||||
modelId: string | string[],
|
||||
engineConfig?: MLCEngineConfig,
|
||||
chatOpts?: ChatOptions | ChatOptions[],
|
||||
keepAliveMs = 10000,
|
||||
): Promise<ServiceWorkerMLCEngine> {
|
||||
if (!("serviceWorker" in navigator)) {
|
||||
throw new NoServiceWorkerAPIError();
|
||||
}
|
||||
const serviceWorkerAPI = navigator.serviceWorker as ServiceWorkerContainer;
|
||||
const registration = await serviceWorkerAPI.ready;
|
||||
const serviceWorker = registration.active || serviceWorkerAPI.controller;
|
||||
if (!serviceWorker) {
|
||||
throw new ServiceWorkerInitializationError();
|
||||
}
|
||||
const serviceWorkerMLCEngine = new ServiceWorkerMLCEngine(
|
||||
engineConfig,
|
||||
keepAliveMs,
|
||||
);
|
||||
await serviceWorkerMLCEngine.reload(modelId, chatOpts);
|
||||
return serviceWorkerMLCEngine;
|
||||
}
|
||||
|
||||
/**
|
||||
* A client of MLCEngine that exposes the same interface
|
||||
*/
|
||||
export class ServiceWorkerMLCEngine extends WebWorkerMLCEngine {
|
||||
missedHeartbeat = 0;
|
||||
|
||||
constructor(engineConfig?: MLCEngineConfig, keepAliveMs = 10000) {
|
||||
if (!("serviceWorker" in navigator)) {
|
||||
throw new NoServiceWorkerAPIError();
|
||||
}
|
||||
super(new ServiceWorker(), engineConfig);
|
||||
|
||||
// Keep alive through periodical heartbeat signals
|
||||
setInterval(() => {
|
||||
this.worker.postMessage({ kind: "keepAlive", uuid: crypto.randomUUID() });
|
||||
this.missedHeartbeat += 1;
|
||||
log.trace("missedHeartbeat", this.missedHeartbeat);
|
||||
}, keepAliveMs);
|
||||
}
|
||||
|
||||
onmessage(event: any): void {
|
||||
const msg = event.data;
|
||||
log.trace(
|
||||
`MLC client message: [${msg.kind}] ${JSON.stringify(msg.content)}`,
|
||||
);
|
||||
try {
|
||||
if (msg.kind === "heartbeat") {
|
||||
this.missedHeartbeat = 0;
|
||||
return;
|
||||
}
|
||||
super.onmessage(msg);
|
||||
} catch (err: any) {
|
||||
// This is expected to throw if user has multiple windows open
|
||||
if (!err.message.startsWith("return from a unknown uuid")) {
|
||||
log.error("CreateWebServiceWorkerMLCEngine.onmessage", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+449
@@ -0,0 +1,449 @@
|
||||
/** Util methods. */
|
||||
import { Tokenizer } from "@mlc-ai/web-tokenizers";
|
||||
import { AppConfig, MessagePlaceholders, ModelRecord } from "./config";
|
||||
import {
|
||||
ChatCompletionChunk,
|
||||
ChatCompletionContentPartImage,
|
||||
ChatCompletionMessageToolCall,
|
||||
} from "./openai_api_protocols/index";
|
||||
import {
|
||||
ModelNotFoundError,
|
||||
ModelNotLoadedError,
|
||||
PrefillChunkSizeSmallerThanImageError,
|
||||
SpecifiedModelNotFoundError,
|
||||
ToolCallOutputInvalidTypeError,
|
||||
ToolCallOutputMissingFieldsError,
|
||||
ToolCallOutputParseError,
|
||||
UnclearModelToUseError,
|
||||
} from "./error";
|
||||
|
||||
/**
|
||||
* Based on `p_prob` of size (vocabSize,) which becomes a distribution after calling
|
||||
* `applySoftmaxWithTemperature()`, sample `top_logprobs` top-probable tokens.
|
||||
*
|
||||
* @param num_top_probs: `top_logprobs` from ChatCompletionRequest
|
||||
* @param p_prob: `logitsOnCPUArray`, being a distribution after `applySoftmaxWithTemperature()`.
|
||||
*
|
||||
* Followed implementation of `ComputeTopProbsImpl()` from [https://github.com/mlc-ai/mlc-llm/blob/
|
||||
* 5b8c529e9704abd09b0432da6dcb4b013fdf43b1/cpp/serve/sampler/cpu_sampler.cc].
|
||||
*
|
||||
* @returns Arrays of (tokenID, prob) pairs, ranked from highest prob to least.
|
||||
*/
|
||||
export function getTopProbs(
|
||||
num_top_probs: number,
|
||||
p_prob: Float32Array,
|
||||
): Array<[number, number]> {
|
||||
if (num_top_probs == 0) return [];
|
||||
// Initialize to dummy values
|
||||
const top_probs: Array<[number, number]> = [];
|
||||
const ndata = p_prob.length;
|
||||
for (let i = 0; i < num_top_probs; i++) {
|
||||
top_probs.push([-1, -1.0]);
|
||||
}
|
||||
|
||||
let sum_prob = 0.0;
|
||||
// Selection argsort.
|
||||
for (let p = 0; p < ndata; p++) {
|
||||
let i = num_top_probs - 1;
|
||||
for (; i >= 0; --i) {
|
||||
if (p_prob[p] > top_probs[i][1]) {
|
||||
if (i !== num_top_probs - 1) {
|
||||
top_probs[i + 1] = top_probs[i];
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (i !== num_top_probs - 1) {
|
||||
top_probs[i + 1] = [p, p_prob[p]];
|
||||
}
|
||||
|
||||
// Early exit
|
||||
sum_prob += p_prob[p];
|
||||
if (1 - sum_prob <= top_probs[num_top_probs - 1][1]) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
return top_probs;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the token table in the form of a string list of tokens, ordered by their token id.
|
||||
* @param tokenizer A loaded tokenizer.
|
||||
* @note The size of the table (i.e. tokenizer.getVocabSize()) may be smaller than the `vocab_size`
|
||||
* in config.json (length of logits), see https://github.com/QwenLM/Qwen2/issues/147 and
|
||||
* https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/discussions/47.
|
||||
*/
|
||||
export function getTokenTableFromTokenizer(tokenizer: Tokenizer): string[] {
|
||||
const tokenTable: string[] = [];
|
||||
const vocabSize = tokenizer.getVocabSize();
|
||||
for (let tokenId = 0; tokenId < vocabSize; tokenId++) {
|
||||
tokenTable.push(tokenizer.idToToken(tokenId));
|
||||
}
|
||||
return tokenTable;
|
||||
}
|
||||
|
||||
/**
|
||||
* Postprocess the suffix of ModelRecord.model to be "/resolve/main/" if it is not specified otherwise.
|
||||
* e.g. https://huggingface.co/mlc-ai/OpenHermes-2.5-Mistral-7B-q4f16_1-MLC/resolve/main/
|
||||
* @return the href of the final URL.
|
||||
*/
|
||||
export function cleanModelUrl(modelUrl: string): string {
|
||||
// https://huggingface.co/USER/MODEL -> https://huggingface.co/USER/MODEL/
|
||||
modelUrl += modelUrl.endsWith("/") ? "" : "/";
|
||||
if (!modelUrl.match(/.+\/resolve\/.+\//)) modelUrl += "resolve/main/";
|
||||
// https://huggingface.co/USER/MODEL/ -> https://huggingface.co/USER/MODEL/resolve/main/
|
||||
return new URL(modelUrl).href;
|
||||
}
|
||||
|
||||
// Constants for Hermes-2-Pro / Hermes-3 models function calling
|
||||
// Follows https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B#prompt-format-for-function-calling
|
||||
// https://huggingface.co/NousResearch/Hermes-3-Llama-3.1-8B#prompt-format-for-function-calling
|
||||
|
||||
/**
|
||||
* Json schema used to prompt the model for function calling; directly copied from the official guide.
|
||||
* This represents to a single function call.
|
||||
*/
|
||||
export const officialHermes2FunctionCallSchema = `{"properties": {"arguments": {"title": "Arguments", "type": "object"}, "name": {"title": "Name", "type": "string"}}, "required": ["arguments", "name"], "title": "FunctionCall", "type": "object"}`;
|
||||
|
||||
/**
|
||||
* A list of such function calls. Used to specify response format, since the output is expected to
|
||||
* be a list of such function calls.
|
||||
*/
|
||||
export const officialHermes2FunctionCallSchemaArray = `{"type":"array","items":${officialHermes2FunctionCallSchema}}`;
|
||||
|
||||
/**
|
||||
* Full system prompt for Hermes-2-Pro and Hermes-3 function calling.
|
||||
*/
|
||||
export const hermes2FunctionCallingSystemPrompt = `You are a function calling AI model. You are
|
||||
provided with function signatures within <tools></tools> XML tags. You may call one or more functions
|
||||
to assist with the user query. Don't make assumptions about what values to plug into functions. Here
|
||||
are the available tools: <tools> ${MessagePlaceholders.hermes_tools} </tools>.
|
||||
Use the following pydantic model json schema for each tool call you will make:
|
||||
${officialHermes2FunctionCallSchema} For each function call return a json object.`;
|
||||
|
||||
/**
|
||||
* Given a string outputMessage, parse it as a JSON object and return an array of tool calls.
|
||||
*
|
||||
* Expect outputMessage to be a valid JSON string, and expect it to be an array of Function with
|
||||
* fields `arguments` and `name`.
|
||||
*/
|
||||
export function getToolCallFromOutputMessage(
|
||||
outputMessage: string,
|
||||
isStreaming: false,
|
||||
): Array<ChatCompletionMessageToolCall>;
|
||||
export function getToolCallFromOutputMessage(
|
||||
outputMessage: string,
|
||||
isStreaming: true,
|
||||
): Array<ChatCompletionChunk.Choice.Delta.ToolCall>;
|
||||
export function getToolCallFromOutputMessage(
|
||||
outputMessage: string,
|
||||
isStreaming: boolean,
|
||||
):
|
||||
| Array<ChatCompletionMessageToolCall>
|
||||
| Array<ChatCompletionChunk.Choice.Delta.ToolCall> {
|
||||
// 1. Parse outputMessage to JSON object
|
||||
let toolCallsObject;
|
||||
try {
|
||||
toolCallsObject = JSON.parse(outputMessage);
|
||||
} catch (err) {
|
||||
throw new ToolCallOutputParseError(outputMessage, err as Error);
|
||||
}
|
||||
|
||||
// 2. Expect to be an array
|
||||
if (!(toolCallsObject instanceof Array)) {
|
||||
throw new ToolCallOutputInvalidTypeError("array");
|
||||
}
|
||||
|
||||
// 3. Parse each tool call and populate tool_calls
|
||||
const numToolCalls = toolCallsObject.length;
|
||||
const tool_calls = [];
|
||||
for (let id = 0; id < numToolCalls; id++) {
|
||||
const curToolCall = toolCallsObject[id];
|
||||
if (curToolCall.name === undefined || curToolCall.arguments === undefined) {
|
||||
throw new ToolCallOutputMissingFieldsError(
|
||||
["name", "arguments"],
|
||||
curToolCall,
|
||||
);
|
||||
}
|
||||
tool_calls.push({
|
||||
name: curToolCall.name,
|
||||
arguments: JSON.stringify(curToolCall.arguments),
|
||||
});
|
||||
}
|
||||
|
||||
// 4. Return based on whether it is streaming or not
|
||||
if (isStreaming) {
|
||||
const tool_calls_result: Array<ChatCompletionChunk.Choice.Delta.ToolCall> =
|
||||
[];
|
||||
for (let id = 0; id < numToolCalls; id++) {
|
||||
const curToolCall = tool_calls[id];
|
||||
tool_calls_result.push({
|
||||
index: id,
|
||||
function: {
|
||||
name: curToolCall.name,
|
||||
arguments: curToolCall.arguments,
|
||||
},
|
||||
type: "function",
|
||||
});
|
||||
}
|
||||
return tool_calls_result;
|
||||
} else {
|
||||
const tool_calls_result: Array<ChatCompletionMessageToolCall> = [];
|
||||
for (let id = 0; id < numToolCalls; id++) {
|
||||
const curToolCall = tool_calls[id];
|
||||
tool_calls_result.push({
|
||||
id: id.toString(),
|
||||
function: {
|
||||
name: curToolCall.name,
|
||||
arguments: curToolCall.arguments,
|
||||
},
|
||||
type: "function",
|
||||
});
|
||||
}
|
||||
return tool_calls_result;
|
||||
}
|
||||
}
|
||||
|
||||
export function findModelRecord(
|
||||
modelId: string,
|
||||
appConfig: AppConfig,
|
||||
): ModelRecord {
|
||||
const matchedItem = appConfig.model_list.find(
|
||||
(item) => item.model_id == modelId,
|
||||
);
|
||||
if (matchedItem !== undefined) return matchedItem;
|
||||
throw new ModelNotFoundError(modelId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the model to use given the loaded modelIds and requestModel. Throws error when unclear
|
||||
* which model to load.
|
||||
* @param loadedModelIds Models currently loaded in the engine.
|
||||
* @param requestModel Model the user specified to load via the request. Required when multiple
|
||||
* models are loaded
|
||||
* @param requestName The type of request or API to load the model for. Needed for error throwing.
|
||||
*/
|
||||
export function getModelIdToUse(
|
||||
loadedModelIds: string[],
|
||||
requestModel: string | undefined | null,
|
||||
requestName: string,
|
||||
): string {
|
||||
let selectedModelId: string;
|
||||
if (loadedModelIds.length === 0) {
|
||||
throw new ModelNotLoadedError(requestName);
|
||||
}
|
||||
if (requestModel) {
|
||||
// If specified model
|
||||
if (loadedModelIds.indexOf(requestModel) === -1) {
|
||||
throw new SpecifiedModelNotFoundError(
|
||||
loadedModelIds,
|
||||
requestModel,
|
||||
requestName,
|
||||
);
|
||||
} else {
|
||||
selectedModelId = requestModel;
|
||||
}
|
||||
} else {
|
||||
// If not specified
|
||||
if (loadedModelIds.length > 1) {
|
||||
throw new UnclearModelToUseError(loadedModelIds, requestName);
|
||||
} else {
|
||||
selectedModelId = loadedModelIds[0];
|
||||
}
|
||||
}
|
||||
return selectedModelId;
|
||||
}
|
||||
|
||||
/**
|
||||
* TODO: Consider if this is the best strategy (though aligned with mlc-llm). We currently greedily
|
||||
* try to fill up prefillChunkSize. Consider the example with 2048 prefill chunk size:
|
||||
* const inputData = [
|
||||
image1, // 1921
|
||||
rangeArr(0, 2048),
|
||||
image2,
|
||||
];
|
||||
* Current approach results in chunks:
|
||||
[image1, rangeArr(0, 127)],
|
||||
[rangeArr(127, 2048)],
|
||||
[image2],
|
||||
* This means 4 embedding kernels and 3 prefill kernels.
|
||||
* While the optimal chunking may be:
|
||||
[image1],
|
||||
[rangeArr(0, 2048)],
|
||||
[image2],
|
||||
* This results in 3 embedding kernels and 3 prefill kernels.
|
||||
* However, greedy strategy is more intuitive and probably more generalizable.
|
||||
*/
|
||||
|
||||
/**
|
||||
* Chunk the inputData such that each chunk's total input length is smaller than prefill
|
||||
* chunk size.
|
||||
* @returns [the data chunks, the input length of each chunk]
|
||||
* @note precondition: if inputData has image in it, then prefillChunkSize >= imageEmbedSize.
|
||||
*/
|
||||
export function getChunkedPrefillInputData(
|
||||
inputData: Array<Array<number> | ImageURL>,
|
||||
prefillChunkSize: number,
|
||||
getImageEmbedSize: (image: ImageURL) => number,
|
||||
): [Array<Array<number> | ImageURL>[], Array<number>] {
|
||||
const chunks: Array<Array<number> | ImageURL>[] = [];
|
||||
const chunkLens: Array<number> = [];
|
||||
let curChunk: Array<Array<number> | ImageURL> = [];
|
||||
let curChunkLen = 0;
|
||||
for (let i = 0; i < inputData.length; i++) {
|
||||
let curData: Array<number> | ImageURL = inputData[i];
|
||||
const curDataLen = Array.isArray(curData)
|
||||
? curData.length
|
||||
: getImageEmbedSize(curData);
|
||||
if (!Array.isArray(curData) && curDataLen > prefillChunkSize) {
|
||||
throw new PrefillChunkSizeSmallerThanImageError(
|
||||
prefillChunkSize,
|
||||
curDataLen,
|
||||
);
|
||||
}
|
||||
// 1. curData can fit into this chunk
|
||||
if (curChunkLen + curDataLen <= prefillChunkSize) {
|
||||
curChunk.push(curData);
|
||||
curChunkLen += curDataLen;
|
||||
if (curChunkLen === prefillChunkSize) {
|
||||
chunks.push([...curChunk]);
|
||||
chunkLens.push(curChunkLen);
|
||||
curChunk = [];
|
||||
curChunkLen = 0;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// 2. Otherwise, depends on whether it is token data or image data
|
||||
if (Array.isArray(curData)) {
|
||||
// 2.1. Token data, which itself needs to be chunked. Keep
|
||||
// chunking and finalizing until finished
|
||||
while (curData.length > 0) {
|
||||
const curDataToChunkLen = Math.min(
|
||||
curData.length,
|
||||
prefillChunkSize - curChunkLen,
|
||||
);
|
||||
curChunk.push(curData.slice(0, curDataToChunkLen));
|
||||
curChunkLen += curDataToChunkLen;
|
||||
curData = curData.slice(curDataToChunkLen);
|
||||
if (curChunkLen === prefillChunkSize) {
|
||||
// curChunk is now full, so finalize to chunks
|
||||
chunks.push([...curChunk]);
|
||||
chunkLens.push(curChunkLen);
|
||||
curChunk = [];
|
||||
curChunkLen = 0;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 2.2. Image data, which itself cannot be chunked, so cannot fit in current chunk.
|
||||
// 2.2.1. Finalize curChunk
|
||||
if (curChunk.length === 0) {
|
||||
throw new Error(
|
||||
"InternalError: do not expect curChunk to be empty when an image does not fit.",
|
||||
);
|
||||
}
|
||||
chunks.push([...curChunk]);
|
||||
chunkLens.push(curChunkLen);
|
||||
// 2.2.2. Then push image to the new chunk
|
||||
curChunk = [curData];
|
||||
curChunkLen = curDataLen;
|
||||
if (curChunkLen === prefillChunkSize) {
|
||||
chunks.push([...curChunk]);
|
||||
chunkLens.push(curChunkLen);
|
||||
curChunk = [];
|
||||
curChunkLen = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Last chunk
|
||||
if (curChunk.length > 0) {
|
||||
chunks.push([...curChunk]);
|
||||
chunkLens.push(curChunkLen);
|
||||
}
|
||||
|
||||
return [chunks, chunkLens];
|
||||
}
|
||||
|
||||
type Cont = () => void;
|
||||
|
||||
/**
|
||||
* A lock implemented using Promise.
|
||||
*
|
||||
* Referred to:
|
||||
* - https://jackpordi.com/posts/locks-in-js-because-why-not
|
||||
* - https://www.linkedin.com/pulse/asynchronous-locking-using-promises-javascript-abdul-ahad-o7smf/
|
||||
*/
|
||||
export class CustomLock {
|
||||
private acquired = false;
|
||||
private readonly queue: Cont[] = [];
|
||||
|
||||
public async acquire(): Promise<void> {
|
||||
if (!this.acquired) {
|
||||
// If lock is free, directly return
|
||||
this.acquired = true;
|
||||
} else {
|
||||
// Otherwise, push the request to the queue, and
|
||||
// a future release() will resolve it
|
||||
return new Promise<void>((resolve) => {
|
||||
this.queue.push(resolve);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
public async release(): Promise<void> {
|
||||
if (!this.acquired) {
|
||||
throw Error("InternalError: expect lock is acquired upon release()");
|
||||
}
|
||||
if (this.queue.length === 0) {
|
||||
// No one is waiting for the lock, so we free it
|
||||
this.acquired = false;
|
||||
return;
|
||||
}
|
||||
|
||||
// Otherwise, hand the execution to the next in queue, and
|
||||
// the lock is still acquired
|
||||
const cont = this.queue.shift();
|
||||
return new Promise((res: Cont) => {
|
||||
cont!();
|
||||
res();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Image related
|
||||
type ImageURL = ChatCompletionContentPartImage.ImageURL;
|
||||
|
||||
/**
|
||||
* Given a url, get the image data. The url can either start with `http` or `data:image`.
|
||||
*/
|
||||
export async function getImageDataFromURL(url: string): Promise<ImageData> {
|
||||
const response = await fetch(url, { mode: "cors" });
|
||||
const img = await createImageBitmap(await response.blob());
|
||||
const canvas = new OffscreenCanvas(img.width, img.height);
|
||||
const ctx = canvas.getContext("2d");
|
||||
if (!ctx) {
|
||||
throw new Error("Could not get 2d context");
|
||||
}
|
||||
ctx.drawImage(img, 0, 0);
|
||||
|
||||
const imageData = ctx.getImageData(0, 0, img.width, img.height);
|
||||
return imageData;
|
||||
}
|
||||
|
||||
/**
|
||||
* Given an ImageData, return the RGB array in Uint8ClampedArray. Note the ImageData.data
|
||||
* is RGBA, so we skip every fourth element of the data. The order goes by rows from the
|
||||
* top-left pixel to the bottom-right, in RGB order.
|
||||
*/
|
||||
export function getRGBArrayFromImageData(
|
||||
imageData: ImageData,
|
||||
): Uint8ClampedArray {
|
||||
const newData = new Uint8ClampedArray(imageData.width * imageData.height * 3);
|
||||
for (let i = 0, offset = 0; i < imageData.data.length; i += 4) {
|
||||
newData[offset++] = imageData.data[i];
|
||||
newData[offset++] = imageData.data[i + 1];
|
||||
newData[offset++] = imageData.data[i + 2];
|
||||
}
|
||||
return newData;
|
||||
}
|
||||
+262
@@ -0,0 +1,262 @@
|
||||
import { AppConfig, ChatOptions } from "./config";
|
||||
import {
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionRequestBase,
|
||||
ChatCompletionRequestStreaming,
|
||||
ChatCompletionRequestNonStreaming,
|
||||
ChatCompletion,
|
||||
ChatCompletionChunk,
|
||||
CompletionCreateParams,
|
||||
Completion,
|
||||
CompletionCreateParamsBase,
|
||||
CompletionCreateParamsStreaming,
|
||||
CompletionCreateParamsNonStreaming,
|
||||
EmbeddingCreateParams,
|
||||
CreateEmbeddingResponse,
|
||||
} from "./openai_api_protocols/index";
|
||||
import * as API from "./openai_api_protocols/index";
|
||||
|
||||
/**
|
||||
* Report during intialization.
|
||||
*/
|
||||
export interface InitProgressReport {
|
||||
progress: number;
|
||||
timeElapsed: number;
|
||||
text: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Callbacks used to report initialization process.
|
||||
*/
|
||||
export type InitProgressCallback = (report: InitProgressReport) => void;
|
||||
|
||||
/**
|
||||
* A stateful logitProcessor used to post-process logits after forwarding the input and before
|
||||
* sampling the next token. If used with `GenerationConfig.logit_bias`, logit_bias is applied after
|
||||
* `processLogits()` is called.
|
||||
*/
|
||||
export interface LogitProcessor {
|
||||
/**
|
||||
* Process logits after forward() and before sampling implicitly, happens on the CPU.
|
||||
* @param logits The logits right after forward().
|
||||
* Returns the processed logits.
|
||||
*/
|
||||
processLogits: (logits: Float32Array) => Float32Array;
|
||||
|
||||
/**
|
||||
* Use the sampled token to update the LogitProcessor's internal state. Called implicitly
|
||||
* right after the next token is sampled/committed.
|
||||
* @param token Token sampled from the processed logits.
|
||||
*/
|
||||
processSampledToken: (token: number) => void;
|
||||
|
||||
/**
|
||||
* Called when in `MLCEngine.resetChat()`. Can clear internal states.
|
||||
*/
|
||||
resetState: () => void;
|
||||
}
|
||||
|
||||
/**
|
||||
* Common interface of MLCEngine that UI can interact with
|
||||
*/
|
||||
export interface MLCEngineInterface {
|
||||
/**
|
||||
* An object that exposes chat-related APIs.
|
||||
*/
|
||||
chat: API.Chat;
|
||||
|
||||
/**
|
||||
* An object that exposes text completion APIs.
|
||||
*/
|
||||
completions: API.Completions;
|
||||
|
||||
/**
|
||||
* An object that exposes embeddings APIs.
|
||||
*/
|
||||
embeddings: API.Embeddings;
|
||||
|
||||
/**
|
||||
* Set an initialization progress callback function
|
||||
* which reports the progress of model loading.
|
||||
*
|
||||
* This function can be useful to implement an UI that
|
||||
* update as we loading the model.
|
||||
*
|
||||
* @param initProgressCallback The callback function
|
||||
*/
|
||||
setInitProgressCallback: (initProgressCallback: InitProgressCallback) => void;
|
||||
|
||||
/**
|
||||
* @returns The current initialization progress callback function.
|
||||
*/
|
||||
getInitProgressCallback: () => InitProgressCallback | undefined;
|
||||
|
||||
/**
|
||||
* Setter for the engine's appConfig.
|
||||
*/
|
||||
setAppConfig: (appConfig: AppConfig) => void;
|
||||
|
||||
/**
|
||||
* Reload the chat with a new model.
|
||||
*
|
||||
* @param modelId model_id of the model to load, either string or string[]. When multiple models
|
||||
* are provided, we load all models sequentially. Each modelId needs to either be in
|
||||
* `webllm.prebuiltAppConfig`, or in `engineConfig.appConfig`.
|
||||
* @param chatOpts Extra options to optionally override the `mlc-chat-config.json` of `modelId`.
|
||||
* The size of which needs to match that of `modelId`; chatOpts[i] will be used for modelId[i].
|
||||
* @returns A promise when reload finishes.
|
||||
* @throws Throws error when device lost (mostly due to OOM); users should re-call reload(),
|
||||
* potentially with a smaller model or smaller context window size.
|
||||
* @note This is an async function.
|
||||
*/
|
||||
reload: (
|
||||
modelId: string | string[],
|
||||
chatOpts?: ChatOptions | ChatOptions[],
|
||||
) => Promise<void>;
|
||||
|
||||
/**
|
||||
* OpenAI-style API. Generate a chat completion response for the given conversation and
|
||||
* configuration. Use `engine.chat.completions.create()` to invoke this API.
|
||||
*
|
||||
* @param request A OpenAI-style ChatCompletion request.
|
||||
*
|
||||
* @note The API is completely functional in behavior. That is, a previous request would not
|
||||
* affect the current request's result. Thus, for multi-round chatting, users are responsible for
|
||||
* maintaining the chat history. With that being said, as an implicit internal optimization, if we
|
||||
* detect that the user is performing multi-round chatting, we will preserve the KV cache and only
|
||||
* prefill the new tokens.
|
||||
* @note For requests sent to the same modelId, will block until all previous requests finish.
|
||||
* @note For more, see https://platform.openai.com/docs/api-reference/chat
|
||||
*/
|
||||
chatCompletion(
|
||||
request: ChatCompletionRequestNonStreaming,
|
||||
): Promise<ChatCompletion>;
|
||||
chatCompletion(
|
||||
request: ChatCompletionRequestStreaming,
|
||||
): Promise<AsyncIterable<ChatCompletionChunk>>;
|
||||
chatCompletion(
|
||||
request: ChatCompletionRequestBase,
|
||||
): Promise<AsyncIterable<ChatCompletionChunk> | ChatCompletion>;
|
||||
chatCompletion(
|
||||
request: ChatCompletionRequest,
|
||||
): Promise<AsyncIterable<ChatCompletionChunk> | ChatCompletion>;
|
||||
|
||||
/**
|
||||
* OpenAI-style API. Completes a CompletionCreateParams, a text completion with no chat template.
|
||||
* Use `engine.completions.create()` to invoke this API.
|
||||
*
|
||||
* @param request An OpenAI-style Completion request.
|
||||
*
|
||||
* @note For requests sent to the same modelId, will block until all previous requests finish.
|
||||
* @note For more, see https://platform.openai.com/docs/api-reference/completions
|
||||
*/
|
||||
completion(request: CompletionCreateParamsNonStreaming): Promise<Completion>;
|
||||
completion(
|
||||
request: CompletionCreateParamsStreaming,
|
||||
): Promise<AsyncIterable<Completion>>;
|
||||
completion(
|
||||
request: CompletionCreateParamsBase,
|
||||
): Promise<AsyncIterable<Completion> | Completion>;
|
||||
completion(
|
||||
request: CompletionCreateParams,
|
||||
): Promise<AsyncIterable<Completion> | Completion>;
|
||||
|
||||
/**
|
||||
* OpenAI-style API. Creates an embedding vector representing the input text.
|
||||
* Use `engine.embeddings.create()` to invoke this API.
|
||||
*
|
||||
* @param request An OpenAI-style Embeddings request.
|
||||
*
|
||||
* @note For requests sent to the same modelId, will block until all previous requests finish.
|
||||
* @note For more, see https://platform.openai.com/docs/api-reference/embeddings/create
|
||||
*/
|
||||
embedding(request: EmbeddingCreateParams): Promise<CreateEmbeddingResponse>;
|
||||
|
||||
/**
|
||||
* @returns A text summarizing the runtime stats.
|
||||
* @param modelId Only required when multiple models are loaded.
|
||||
* @note This is an async function
|
||||
*/
|
||||
runtimeStatsText: (modelId?: string) => Promise<string>;
|
||||
|
||||
/**
|
||||
* Interrupt the generate process if it is already running.
|
||||
*/
|
||||
interruptGenerate: () => void;
|
||||
|
||||
/**
|
||||
* Explicitly unload the currently loaded model(s) and release the related resources. Waits until
|
||||
* the webgpu device finishes all submitted work and destroys itself.
|
||||
* @note This is an asynchronous function.
|
||||
*/
|
||||
unload: () => Promise<void>;
|
||||
|
||||
/**
|
||||
* Reset the current chat session by clear all memories.
|
||||
* @param keepStats: If True, do not reset the statistics.
|
||||
* @param modelId Only required when multiple models are loaded.
|
||||
*/
|
||||
resetChat: (keepStats?: boolean, modelId?: string) => Promise<void>;
|
||||
|
||||
/**
|
||||
* Get the current generated response.
|
||||
* @param modelId Only required when multiple models are loaded.
|
||||
* @returns The current output message.
|
||||
*/
|
||||
getMessage: (modelId?: string) => Promise<string>;
|
||||
|
||||
/**
|
||||
* Returns the device's maxStorageBufferBindingSize, can be used to guess whether the device
|
||||
* has limited resources like an Android phone.
|
||||
*/
|
||||
getMaxStorageBufferBindingSize(): Promise<number>;
|
||||
|
||||
/**
|
||||
* Returns the device's gpu vendor (e.g. arm, qualcomm, apple) if available. Otherwise return
|
||||
* an empty string.
|
||||
*/
|
||||
getGPUVendor(): Promise<string>;
|
||||
|
||||
/**
|
||||
* Forward the given input tokens to the model, then sample the next token.
|
||||
*
|
||||
* This function has side effects as the model will update its KV cache.
|
||||
*
|
||||
* @param inputIds The input tokens.
|
||||
* @param isPrefill True if prefill, false if decode; only used for statistics.
|
||||
* @param modelId Only required when multiple models are loaded.
|
||||
* @returns Next token sampled.
|
||||
* @note This is an async function.
|
||||
*/
|
||||
forwardTokensAndSample(
|
||||
inputIds: Array<number>,
|
||||
isPrefill: boolean,
|
||||
modelId?: string,
|
||||
): Promise<number>;
|
||||
|
||||
/**
|
||||
* Set MLCEngine logging output level
|
||||
*
|
||||
* @param logLevel The new log level
|
||||
*/
|
||||
setLogLevel(logLevel: LogLevel): void;
|
||||
}
|
||||
|
||||
export const LOG_LEVELS = {
|
||||
TRACE: 0,
|
||||
DEBUG: 1,
|
||||
INFO: 2,
|
||||
WARN: 3,
|
||||
ERROR: 4,
|
||||
SILENT: 5,
|
||||
};
|
||||
export type LogLevel = keyof typeof LOG_LEVELS;
|
||||
|
||||
export type LatencyBreakdown = {
|
||||
logitProcessorTime: number[];
|
||||
logitBiasTime: number[];
|
||||
penaltyTime: number[];
|
||||
sampleTime: number[];
|
||||
totalTime: number[];
|
||||
grammarBitmaskTime: number[];
|
||||
};
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
import { AppConfig, ChatOptions, ModelRecord, getCacheBackend } from "./config";
|
||||
|
||||
// Helper function to compare two arrays
|
||||
export function areArraysEqual(arr1?: Array<any>, arr2?: Array<any>): boolean {
|
||||
if (!arr1 && !arr2) return true;
|
||||
if (!arr1 || !arr2) return false;
|
||||
if (arr1.length !== arr2.length) return false;
|
||||
for (let i = 0; i < arr1.length; i++) {
|
||||
if (arr1[i] !== arr2[i]) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// Helper function to compare two objects deeply
|
||||
function areObjectsEqual(obj1: any, obj2: any): boolean {
|
||||
if (obj1 === obj2) return true;
|
||||
if (typeof obj1 !== typeof obj2) return false;
|
||||
if (typeof obj1 !== "object" || obj1 === null || obj2 === null) return false;
|
||||
|
||||
const keys1 = Object.keys(obj1);
|
||||
const keys2 = Object.keys(obj2);
|
||||
if (keys1.length !== keys2.length) return false;
|
||||
|
||||
for (const key of keys1) {
|
||||
if (!keys2.includes(key) || !areObjectsEqual(obj1[key], obj2[key]))
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// Function to compare two ModelRecord instances
|
||||
export function areModelRecordsEqual(
|
||||
record1: ModelRecord,
|
||||
record2: ModelRecord,
|
||||
): boolean {
|
||||
// Compare primitive fields
|
||||
if (
|
||||
record1.model !== record2.model ||
|
||||
record1.model_id !== record2.model_id ||
|
||||
record1.model_lib !== record2.model_lib ||
|
||||
record1.vram_required_MB !== record2.vram_required_MB ||
|
||||
record1.low_resource_required !== record2.low_resource_required ||
|
||||
record1.buffer_size_required_bytes !== record2.buffer_size_required_bytes
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Compare required_features arrays
|
||||
if (
|
||||
(record1.required_features && !record2.required_features) ||
|
||||
(!record1.required_features && record2.required_features)
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (record1.required_features && record2.required_features) {
|
||||
if (record1.required_features.length !== record2.required_features.length) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (let i = 0; i < record1.required_features.length; i++) {
|
||||
if (record1.required_features[i] !== record2.required_features[i]) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
export function areAppConfigsEqual(
|
||||
config1?: AppConfig,
|
||||
config2?: AppConfig,
|
||||
): boolean {
|
||||
if (config1 === undefined || config2 === undefined) {
|
||||
return config1 === config2;
|
||||
}
|
||||
|
||||
// Check if both configurations have the same cache backend
|
||||
if (getCacheBackend(config1) !== getCacheBackend(config2)) {
|
||||
return false;
|
||||
}
|
||||
if (
|
||||
getCacheBackend(config1) === "opfs" &&
|
||||
(config1.opfsAccessMode ?? "async") !== (config2.opfsAccessMode ?? "async")
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check if both configurations have the same number of model records
|
||||
if (config1.model_list.length !== config2.model_list.length) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Compare each ModelRecord in the model_list
|
||||
for (let i = 0; i < config1.model_list.length; i++) {
|
||||
if (!areModelRecordsEqual(config1.model_list[i], config2.model_list[i])) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// If all checks passed, the configurations are equal
|
||||
return true;
|
||||
}
|
||||
|
||||
export function areChatOptionsEqual(
|
||||
options1?: ChatOptions,
|
||||
options2?: ChatOptions,
|
||||
): boolean {
|
||||
if (options1 === undefined || options2 === undefined) {
|
||||
return options1 === options2;
|
||||
}
|
||||
// Compare each property of ChatOptions (which are Partial<ChatConfig>)
|
||||
if (!areArraysEqual(options1.tokenizer_files, options2.tokenizer_files))
|
||||
return false;
|
||||
if (!areObjectsEqual(options1.conv_config, options2.conv_config))
|
||||
return false;
|
||||
if (options1.conv_template !== options2.conv_template) return false;
|
||||
if (options1.repetition_penalty !== options2.repetition_penalty) return false;
|
||||
if (options1.frequency_penalty !== options2.frequency_penalty) return false;
|
||||
if (options1.presence_penalty !== options2.presence_penalty) return false;
|
||||
if (options1.top_p !== options2.top_p) return false;
|
||||
if (options1.temperature !== options2.temperature) return false;
|
||||
if (options1.bos_token_id !== options2.bos_token_id) return false;
|
||||
|
||||
// If all checks passed, the options are equal
|
||||
return true;
|
||||
}
|
||||
|
||||
export function areChatOptionsListEqual(
|
||||
options1?: ChatOptions[],
|
||||
options2?: ChatOptions[],
|
||||
): boolean {
|
||||
if (options1 && options2) {
|
||||
// Both defined, need to compare
|
||||
if (options1.length !== options2.length) {
|
||||
return false;
|
||||
} else {
|
||||
for (let i = 0; i < options1.length; i++) {
|
||||
if (!areChatOptionsEqual(options1[i], options2[i])) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
} else if (!options1 && !options2) {
|
||||
// Both undefined, equal
|
||||
return true;
|
||||
} else {
|
||||
// One defined, other not
|
||||
return false;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,842 @@
|
||||
import { AppConfig, ChatOptions, MLCEngineConfig } from "./config";
|
||||
import {
|
||||
MLCEngineInterface,
|
||||
InitProgressCallback,
|
||||
InitProgressReport,
|
||||
LogLevel,
|
||||
LogitProcessor,
|
||||
} from "./types";
|
||||
import {
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionRequestBase,
|
||||
ChatCompletionRequestStreaming,
|
||||
ChatCompletionRequestNonStreaming,
|
||||
ChatCompletion,
|
||||
ChatCompletionChunk,
|
||||
Completion,
|
||||
CompletionCreateParamsNonStreaming,
|
||||
CompletionCreateParamsStreaming,
|
||||
CompletionCreateParamsBase,
|
||||
CompletionCreateParams,
|
||||
CreateEmbeddingResponse,
|
||||
EmbeddingCreateParams,
|
||||
} from "./openai_api_protocols/index";
|
||||
import * as API from "./openai_api_protocols/index";
|
||||
import {
|
||||
MessageContent,
|
||||
ReloadParams,
|
||||
ForwardTokensAndSampleParams,
|
||||
ChatCompletionNonStreamingParams,
|
||||
ChatCompletionStreamInitParams,
|
||||
ResetChatParams,
|
||||
WorkerResponse,
|
||||
WorkerRequest,
|
||||
CompletionNonStreamingParams,
|
||||
EmbeddingParams,
|
||||
CompletionStreamInitParams,
|
||||
GetMessageParams,
|
||||
RuntimeStatsTextParams,
|
||||
CompletionStreamNextChunkParams,
|
||||
} from "./message";
|
||||
import log from "loglevel";
|
||||
import { MLCEngine } from "./engine";
|
||||
import {
|
||||
UnknownMessageKindError,
|
||||
WorkerEngineModelNotLoadedError,
|
||||
} from "./error";
|
||||
import { areArraysEqual } from "./utils";
|
||||
import { getModelIdToUse } from "./support";
|
||||
|
||||
/**
|
||||
* Worker handler that can be used in a WebWorker
|
||||
*
|
||||
* @example
|
||||
*
|
||||
* // setup a chat worker handler that routes
|
||||
* // requests to the chat
|
||||
* const engine = new MLCEngine();
|
||||
* cont handler = new WebWorkerMLCEngineHandler(engine);
|
||||
* onmessage = handler.onmessage;
|
||||
*/
|
||||
export class WebWorkerMLCEngineHandler {
|
||||
/**
|
||||
* The modelId and chatOpts that the underlying engine (backend) is currently loaded with.
|
||||
* An engine can be loaded with multiple models, so modelId and chatOpts are lists.
|
||||
*
|
||||
* TODO(webllm-team): This is always in-sync with `this.engine` unless device is lost due to
|
||||
* unexpected reason. Therefore, we should get it from `this.engine` directly and make handler
|
||||
* stateless. Besides, consider if we should add appConfig, or use engine's API to find the
|
||||
* corresponding model record rather than relying on just the modelId.
|
||||
*/
|
||||
modelId?: string[];
|
||||
chatOpts?: ChatOptions[];
|
||||
|
||||
public engine: MLCEngine;
|
||||
/** ChatCompletion and Completion share the same chunk generator. Each loaded model has its own. */
|
||||
protected loadedModelIdToAsyncGenerator: Map<
|
||||
string,
|
||||
AsyncGenerator<ChatCompletionChunk | Completion, void, void>
|
||||
>;
|
||||
|
||||
/**
|
||||
* @param engine A concrete implementation of MLCEngineInterface
|
||||
*/
|
||||
constructor() {
|
||||
this.engine = new MLCEngine();
|
||||
this.loadedModelIdToAsyncGenerator = new Map<
|
||||
string,
|
||||
AsyncGenerator<ChatCompletionChunk | Completion, void, void>
|
||||
>();
|
||||
this.engine.setInitProgressCallback((report: InitProgressReport) => {
|
||||
const msg: WorkerResponse = {
|
||||
kind: "initProgressCallback",
|
||||
uuid: "",
|
||||
content: report,
|
||||
};
|
||||
this.postMessage(msg);
|
||||
});
|
||||
}
|
||||
|
||||
postMessage(msg: any) {
|
||||
// Use Web Worker DOM Message API
|
||||
postMessage(msg);
|
||||
}
|
||||
|
||||
setLogitProcessorRegistry(
|
||||
logitProcessorRegistry?: Map<string, LogitProcessor>,
|
||||
) {
|
||||
this.engine.setLogitProcessorRegistry(logitProcessorRegistry);
|
||||
}
|
||||
|
||||
async handleTask<T extends MessageContent>(
|
||||
uuid: string,
|
||||
task: () => Promise<T>,
|
||||
) {
|
||||
try {
|
||||
const res = await task();
|
||||
const msg: WorkerResponse = {
|
||||
kind: "return",
|
||||
uuid: uuid,
|
||||
content: res,
|
||||
};
|
||||
this.postMessage(msg);
|
||||
} catch (err) {
|
||||
const errStr = (err as object).toString();
|
||||
const msg: WorkerResponse = {
|
||||
kind: "throw",
|
||||
uuid: uuid,
|
||||
content: errStr,
|
||||
};
|
||||
this.postMessage(msg);
|
||||
}
|
||||
}
|
||||
|
||||
onmessage(
|
||||
event: any,
|
||||
onComplete?: (value: any) => void,
|
||||
onError?: () => void,
|
||||
) {
|
||||
let msg: WorkerRequest;
|
||||
if (event instanceof MessageEvent) {
|
||||
msg = event.data as WorkerRequest;
|
||||
} else {
|
||||
msg = event as WorkerRequest;
|
||||
}
|
||||
switch (msg.kind) {
|
||||
case "reload": {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as ReloadParams;
|
||||
await this.engine.reload(params.modelId, params.chatOpts);
|
||||
this.modelId = params.modelId;
|
||||
this.chatOpts = params.chatOpts;
|
||||
onComplete?.(null);
|
||||
return null;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "forwardTokensAndSample": {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as ForwardTokensAndSampleParams;
|
||||
const res = await this.engine.forwardTokensAndSample(
|
||||
params.inputIds,
|
||||
params.isPrefill,
|
||||
params.modelId,
|
||||
);
|
||||
onComplete?.(res);
|
||||
return res;
|
||||
});
|
||||
return;
|
||||
}
|
||||
// For engine.chat.completions.create()
|
||||
case "chatCompletionNonStreaming": {
|
||||
// Directly return the ChatCompletion response
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as ChatCompletionNonStreamingParams;
|
||||
await this.reloadIfUnmatched(params.modelId, params.chatOpts);
|
||||
const res = await this.engine.chatCompletion(params.request);
|
||||
onComplete?.(res);
|
||||
return res;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "chatCompletionStreamInit": {
|
||||
// One-time set up that instantiates the chunk generator in worker
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as ChatCompletionStreamInitParams;
|
||||
// Also ensures params.selectedModelId will match what this.engine selects
|
||||
await this.reloadIfUnmatched(params.modelId, params.chatOpts);
|
||||
// Register new async generator for this new request of the model
|
||||
const curGenerator = (await this.engine.chatCompletion(
|
||||
params.request,
|
||||
)) as AsyncGenerator<ChatCompletionChunk, void, void>;
|
||||
this.loadedModelIdToAsyncGenerator.set(
|
||||
params.selectedModelId,
|
||||
curGenerator,
|
||||
);
|
||||
onComplete?.(null);
|
||||
return null;
|
||||
});
|
||||
return;
|
||||
}
|
||||
// For engine.completions.create()
|
||||
case "completionNonStreaming": {
|
||||
// Directly return the ChatCompletion response
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as CompletionNonStreamingParams;
|
||||
await this.reloadIfUnmatched(params.modelId, params.chatOpts);
|
||||
const res = await this.engine.completion(params.request);
|
||||
onComplete?.(res);
|
||||
return res;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "completionStreamInit": {
|
||||
// One-time set up that instantiates the chunk generator in worker
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as CompletionStreamInitParams;
|
||||
// Also ensures params.selectedModelId will match what this.engine selects
|
||||
await this.reloadIfUnmatched(params.modelId, params.chatOpts);
|
||||
// Register new async generator for this new request of the model
|
||||
const curGenerator = (await this.engine.completion(
|
||||
params.request,
|
||||
)) as AsyncGenerator<Completion, void, void>;
|
||||
this.loadedModelIdToAsyncGenerator.set(
|
||||
params.selectedModelId,
|
||||
curGenerator,
|
||||
);
|
||||
onComplete?.(null);
|
||||
return null;
|
||||
});
|
||||
return;
|
||||
}
|
||||
// Shared by engine.chat.completions.create() and engine.completions.create()
|
||||
case "completionStreamNextChunk": {
|
||||
// Note: ChatCompletion and Completion share the same chunk generator.
|
||||
// For any subsequent request, we return whatever `next()` yields
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as CompletionStreamNextChunkParams;
|
||||
const curGenerator = this.loadedModelIdToAsyncGenerator.get(
|
||||
params.selectedModelId,
|
||||
);
|
||||
if (curGenerator === undefined) {
|
||||
throw Error(
|
||||
"InternalError: Chunk generator in worker should be instantiated by now.",
|
||||
);
|
||||
}
|
||||
// Yield the next chunk
|
||||
const { value } = await curGenerator.next();
|
||||
onComplete?.(value);
|
||||
return value;
|
||||
});
|
||||
return;
|
||||
}
|
||||
// For engine.embeddings.create()
|
||||
case "embedding": {
|
||||
// Directly return the Embeddings response
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as EmbeddingParams;
|
||||
await this.reloadIfUnmatched(params.modelId, params.chatOpts);
|
||||
const res = await this.engine.embedding(params.request);
|
||||
onComplete?.(res);
|
||||
return res;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "runtimeStatsText": {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as RuntimeStatsTextParams;
|
||||
const res = await this.engine.runtimeStatsText(params.modelId);
|
||||
onComplete?.(res);
|
||||
return res;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "interruptGenerate": {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
this.engine.interruptGenerate();
|
||||
onComplete?.(null);
|
||||
return null;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "unload": {
|
||||
// Unset modelId and chatOpts since backend unloads the model
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
await this.engine.unload();
|
||||
this.modelId = undefined;
|
||||
this.chatOpts = undefined;
|
||||
// This may not be cleaned properly when one asyncGenerator finishes.
|
||||
// We only clear at unload(), which may not be called upon reload().
|
||||
// However, service_worker may skip reload(). Will leave as is for now.
|
||||
this.loadedModelIdToAsyncGenerator.clear();
|
||||
onComplete?.(null);
|
||||
return null;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "resetChat": {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as ResetChatParams;
|
||||
await this.engine.resetChat(params.keepStats, params.modelId);
|
||||
onComplete?.(null);
|
||||
return null;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "getMaxStorageBufferBindingSize": {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const res = await this.engine.getMaxStorageBufferBindingSize();
|
||||
onComplete?.(res);
|
||||
return res;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "getGPUVendor": {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const res = await this.engine.getGPUVendor();
|
||||
onComplete?.(res);
|
||||
return res;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "getMessage": {
|
||||
this.handleTask(msg.uuid, async () => {
|
||||
const params = msg.content as GetMessageParams;
|
||||
const res = await this.engine.getMessage(params.modelId);
|
||||
onComplete?.(res);
|
||||
return res;
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "setLogLevel": {
|
||||
const logLevel = msg.content as LogLevel;
|
||||
this.engine.setLogLevel(logLevel);
|
||||
log.setLevel(logLevel);
|
||||
onComplete?.(null);
|
||||
return;
|
||||
}
|
||||
case "setAppConfig": {
|
||||
const appConfig = msg.content as AppConfig;
|
||||
this.engine.setAppConfig(appConfig);
|
||||
onComplete?.(null);
|
||||
return;
|
||||
}
|
||||
case "customRequest": {
|
||||
onComplete?.(null);
|
||||
return;
|
||||
}
|
||||
default: {
|
||||
if (msg.kind && msg.content) {
|
||||
onError?.();
|
||||
throw new UnknownMessageKindError(msg.kind, msg.content);
|
||||
} else {
|
||||
// Ignore irrelavent events
|
||||
onComplete?.(null);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Check whether frontend expectation matches with backend (modelId and chatOpts). If not (due
|
||||
* to possibly killed service worker), we reload here.
|
||||
* For more, see https://github.com/mlc-ai/web-llm/pull/533
|
||||
*/
|
||||
async reloadIfUnmatched(
|
||||
expectedModelId: string[],
|
||||
expectedChatOpts?: ChatOptions[],
|
||||
) {
|
||||
// TODO: should we also check expectedChatOpts here?
|
||||
if (!areArraysEqual(this.modelId, expectedModelId)) {
|
||||
log.warn(
|
||||
"WebWorkerMLCEngine expects model is loaded in WebWorkerMLCEngineHandler, " +
|
||||
"but it is not. This may due to web/service worker is unexpectedly killed.\n" +
|
||||
"Reloading engine in WebWorkerMLCEngineHandler.",
|
||||
);
|
||||
await this.engine.reload(expectedModelId, expectedChatOpts);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export interface ChatWorker {
|
||||
onmessage: any;
|
||||
postMessage: (message: any) => void;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates `WebWorkerMLCEngine`, a client that holds the same interface as `MLCEngine`.
|
||||
*
|
||||
* Equivalent to `new webllm.WebWorkerMLCEngine(worker).reload(...)`.
|
||||
*
|
||||
* @param worker The worker that holds the actual MLCEngine, initialized with `new Worker()`.
|
||||
* @param modelId model_id of the model to load, either string or string[]. When multiple models
|
||||
* are provided, we load all models sequentially. Each modelId needs to either be in
|
||||
* `webllm.prebuiltAppConfig`, or in `engineCOnfig.appConfig`.
|
||||
* @param engineConfig Optionally configures the engine, see `webllm.MLCEngineConfig` for more.
|
||||
* @param chatOpts Extra options to optionally override the `mlc-chat-config.json` of `modelId`.
|
||||
* The size of which needs to match that of `modelId`; chatOpts[i] will be used for modelId[i].
|
||||
* @returns An initialized `WebLLM.WebWorkerMLCEngine` with `modelId` loaded.
|
||||
*
|
||||
* @note engineConfig.logitProcessorRegistry is ignored for `CreateWebWorkMLCEngine()`.
|
||||
*/
|
||||
export async function CreateWebWorkerMLCEngine(
|
||||
worker: any,
|
||||
modelId: string | string[],
|
||||
engineConfig?: MLCEngineConfig,
|
||||
chatOpts?: ChatOptions | ChatOptions[],
|
||||
): Promise<WebWorkerMLCEngine> {
|
||||
const webWorkerMLCEngine = new WebWorkerMLCEngine(worker, engineConfig);
|
||||
await webWorkerMLCEngine.reload(modelId, chatOpts);
|
||||
return webWorkerMLCEngine;
|
||||
}
|
||||
|
||||
/**
|
||||
* A client of MLCEngine that exposes the same interface
|
||||
*
|
||||
* @example
|
||||
*
|
||||
* const chat = new webllm.WebWorkerMLCEngine(new Worker(
|
||||
* new URL('./worker.ts', import.meta.url),
|
||||
* {type: 'module'}
|
||||
* ));
|
||||
*/
|
||||
export class WebWorkerMLCEngine implements MLCEngineInterface {
|
||||
public worker: ChatWorker;
|
||||
/** For chat.completions.create() */
|
||||
public chat: API.Chat;
|
||||
/** For completions.create() */
|
||||
public completions: API.Completions;
|
||||
/** For embeddings.create() */
|
||||
public embeddings: API.Embeddings;
|
||||
|
||||
/**
|
||||
* The modelId and chatOpts that the frontend expects the backend engine is currently loaded
|
||||
* with. Needed for service worker. It is the backend and handler's job to match up with the
|
||||
* expectation despite the web/service worker possibly being killed.
|
||||
* Since an engine can load multiple models, both modelId and chatOpts are lists.
|
||||
*/
|
||||
modelId?: string[];
|
||||
chatOpts?: ChatOptions[];
|
||||
|
||||
private initProgressCallback?: InitProgressCallback;
|
||||
private pendingPromise = new Map<string, (msg: WorkerResponse) => void>();
|
||||
|
||||
constructor(worker: ChatWorker, engineConfig?: MLCEngineConfig) {
|
||||
this.worker = worker;
|
||||
worker.onmessage = (event: any) => {
|
||||
this.onmessage.bind(this)(event);
|
||||
};
|
||||
|
||||
if (engineConfig?.appConfig) {
|
||||
this.setAppConfig(engineConfig?.appConfig);
|
||||
}
|
||||
if (engineConfig?.logLevel) {
|
||||
this.setLogLevel(engineConfig?.logLevel);
|
||||
}
|
||||
this.setInitProgressCallback(engineConfig?.initProgressCallback);
|
||||
if (engineConfig?.logitProcessorRegistry) {
|
||||
if (engineConfig?.logitProcessorRegistry) {
|
||||
log.warn(
|
||||
"Warning: The `logitProcessorRegistry` property in `engineConfig` will be ignored when using the WebWorkerMLCEngine constructor. To set `logitProcessorRegistry`, use the engine constructor within the worker script instead.",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
this.chat = new API.Chat(this);
|
||||
this.completions = new API.Completions(this);
|
||||
this.embeddings = new API.Embeddings(this);
|
||||
}
|
||||
|
||||
setInitProgressCallback(initProgressCallback?: InitProgressCallback) {
|
||||
this.initProgressCallback = initProgressCallback;
|
||||
}
|
||||
|
||||
getInitProgressCallback(): InitProgressCallback | undefined {
|
||||
return this.initProgressCallback;
|
||||
}
|
||||
|
||||
setAppConfig(appConfig: AppConfig) {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "setAppConfig",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: appConfig,
|
||||
};
|
||||
this.worker.postMessage(msg);
|
||||
}
|
||||
|
||||
setLogLevel(logLevel: LogLevel) {
|
||||
log.setLevel(logLevel);
|
||||
const msg: WorkerRequest = {
|
||||
kind: "setLogLevel",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: logLevel,
|
||||
};
|
||||
this.worker.postMessage(msg);
|
||||
}
|
||||
|
||||
protected getPromise<T extends MessageContent>(
|
||||
msg: WorkerRequest,
|
||||
): Promise<T> {
|
||||
const uuid = msg.uuid;
|
||||
const executor = (
|
||||
resolve: (arg: T) => void,
|
||||
reject: (arg: any) => void,
|
||||
) => {
|
||||
const cb = (msg: WorkerResponse) => {
|
||||
if (msg.kind == "return") {
|
||||
resolve(msg.content as T);
|
||||
} else {
|
||||
if (msg.kind != "throw") {
|
||||
reject("Uknown msg kind " + msg.kind);
|
||||
} else {
|
||||
reject(msg.content);
|
||||
}
|
||||
}
|
||||
};
|
||||
this.pendingPromise.set(uuid, cb);
|
||||
};
|
||||
const promise = new Promise<T>(executor);
|
||||
this.worker.postMessage(msg);
|
||||
return promise;
|
||||
}
|
||||
|
||||
async reload(
|
||||
modelId: string | string[],
|
||||
chatOpts?: ChatOptions | ChatOptions[],
|
||||
): Promise<void> {
|
||||
// Always convert modelId and chatOpts to lists internally for ease of manipulation
|
||||
if (!Array.isArray(modelId)) {
|
||||
modelId = [modelId];
|
||||
}
|
||||
if (chatOpts !== undefined && !Array.isArray(chatOpts)) {
|
||||
chatOpts = [chatOpts];
|
||||
}
|
||||
|
||||
const msg: WorkerRequest = {
|
||||
kind: "reload",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
modelId: modelId,
|
||||
chatOpts: chatOpts,
|
||||
},
|
||||
};
|
||||
await this.getPromise<null>(msg);
|
||||
this.modelId = modelId;
|
||||
this.chatOpts = chatOpts;
|
||||
}
|
||||
|
||||
async getMaxStorageBufferBindingSize(): Promise<number> {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "getMaxStorageBufferBindingSize",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: null,
|
||||
};
|
||||
return await this.getPromise<number>(msg);
|
||||
}
|
||||
|
||||
async getGPUVendor(): Promise<string> {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "getGPUVendor",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: null,
|
||||
};
|
||||
return await this.getPromise<string>(msg);
|
||||
}
|
||||
|
||||
async getMessage(modelId?: string): Promise<string> {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "getMessage",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
modelId: modelId,
|
||||
},
|
||||
};
|
||||
return await this.getPromise<string>(msg);
|
||||
}
|
||||
|
||||
async runtimeStatsText(modelId?: string): Promise<string> {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "runtimeStatsText",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
modelId: modelId,
|
||||
},
|
||||
};
|
||||
return await this.getPromise<string>(msg);
|
||||
}
|
||||
|
||||
interruptGenerate(): void {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "interruptGenerate",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: null,
|
||||
};
|
||||
this.getPromise<null>(msg);
|
||||
}
|
||||
|
||||
async unload(): Promise<void> {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "unload",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: null,
|
||||
};
|
||||
await this.getPromise<null>(msg);
|
||||
this.modelId = undefined;
|
||||
this.chatOpts = undefined;
|
||||
}
|
||||
|
||||
async resetChat(keepStats = false, modelId?: string): Promise<void> {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "resetChat",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
keepStats: keepStats,
|
||||
modelId: modelId,
|
||||
},
|
||||
};
|
||||
await this.getPromise<null>(msg);
|
||||
}
|
||||
|
||||
async forwardTokensAndSample(
|
||||
inputIds: Array<number>,
|
||||
isPrefill: boolean,
|
||||
modelId?: string,
|
||||
): Promise<number> {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "forwardTokensAndSample",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
inputIds: inputIds,
|
||||
isPrefill: isPrefill,
|
||||
modelId: modelId,
|
||||
},
|
||||
};
|
||||
return await this.getPromise<number>(msg);
|
||||
}
|
||||
|
||||
/**
|
||||
* Every time the generator is called, we post a message to the worker asking it to
|
||||
* decode one step, and we expect to receive a message of `ChatCompletionChunk` from
|
||||
* the worker which we yield. The last message is `void`, meaning the generator has nothing
|
||||
* to yield anymore.
|
||||
*
|
||||
* @param selectedModelId: The model of whose async generator to call next() to get next chunk.
|
||||
* Needed because an engine can load multiple models.
|
||||
*
|
||||
* @note ChatCompletion and Completion share the same chunk generator.
|
||||
*/
|
||||
async *asyncGenerate(
|
||||
selectedModelId: string,
|
||||
): AsyncGenerator<ChatCompletionChunk | Completion, void, void> {
|
||||
// Every time it gets called, sends message to worker, asking for the next chunk
|
||||
while (true) {
|
||||
const msg: WorkerRequest = {
|
||||
kind: "completionStreamNextChunk",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
selectedModelId: selectedModelId,
|
||||
} as CompletionStreamNextChunkParams,
|
||||
};
|
||||
const ret = await this.getPromise<ChatCompletionChunk>(msg);
|
||||
// If the worker's generator reached the end, it would return a `void`
|
||||
if (typeof ret !== "object") {
|
||||
break;
|
||||
}
|
||||
yield ret;
|
||||
}
|
||||
}
|
||||
|
||||
async chatCompletion(
|
||||
request: ChatCompletionRequestNonStreaming,
|
||||
): Promise<ChatCompletion>;
|
||||
async chatCompletion(
|
||||
request: ChatCompletionRequestStreaming,
|
||||
): Promise<AsyncIterable<ChatCompletionChunk>>;
|
||||
async chatCompletion(
|
||||
request: ChatCompletionRequestBase,
|
||||
): Promise<AsyncIterable<ChatCompletionChunk> | ChatCompletion>;
|
||||
async chatCompletion(
|
||||
request: ChatCompletionRequest,
|
||||
): Promise<AsyncIterable<ChatCompletionChunk> | ChatCompletion> {
|
||||
if (this.modelId === undefined) {
|
||||
throw new WorkerEngineModelNotLoadedError(this.constructor.name);
|
||||
}
|
||||
// Needed for the streaming case. Consolidate model id to specify
|
||||
// which model's asyncGenerator to instantiate or call next() on.
|
||||
// Since handler can maintain multiple generators concurrently
|
||||
const selectedModelId = getModelIdToUse(
|
||||
this.modelId ? this.modelId : [],
|
||||
request.model,
|
||||
"ChatCompletionRequest",
|
||||
);
|
||||
|
||||
if (request.stream) {
|
||||
// First let worker instantiate a generator
|
||||
const msg: WorkerRequest = {
|
||||
kind: "chatCompletionStreamInit",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
request: request,
|
||||
selectedModelId: selectedModelId,
|
||||
modelId: this.modelId,
|
||||
chatOpts: this.chatOpts,
|
||||
},
|
||||
};
|
||||
await this.getPromise<null>(msg);
|
||||
|
||||
// Then return an async chunk generator that resides on the client side
|
||||
return this.asyncGenerate(selectedModelId) as AsyncGenerator<
|
||||
ChatCompletionChunk,
|
||||
void,
|
||||
void
|
||||
>;
|
||||
}
|
||||
|
||||
// Non streaming case is more straightforward
|
||||
const msg: WorkerRequest = {
|
||||
kind: "chatCompletionNonStreaming",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
request: request,
|
||||
modelId: this.modelId,
|
||||
chatOpts: this.chatOpts,
|
||||
},
|
||||
};
|
||||
return await this.getPromise<ChatCompletion>(msg);
|
||||
}
|
||||
|
||||
async completion(
|
||||
request: CompletionCreateParamsNonStreaming,
|
||||
): Promise<Completion>;
|
||||
async completion(
|
||||
request: CompletionCreateParamsStreaming,
|
||||
): Promise<AsyncIterable<Completion>>;
|
||||
async completion(
|
||||
request: CompletionCreateParamsBase,
|
||||
): Promise<AsyncIterable<Completion> | Completion>;
|
||||
async completion(
|
||||
request: CompletionCreateParams,
|
||||
): Promise<AsyncIterable<Completion> | Completion> {
|
||||
if (this.modelId === undefined) {
|
||||
throw new WorkerEngineModelNotLoadedError(this.constructor.name);
|
||||
}
|
||||
// Needed for the streaming case. Consolidate model id to specify
|
||||
// which model's asyncGenerator to instantiate or call next() on.
|
||||
// Since handler can maintain multiple generators concurrently
|
||||
const selectedModelId = getModelIdToUse(
|
||||
this.modelId ? this.modelId : [],
|
||||
request.model,
|
||||
"CompletionCreateParams",
|
||||
);
|
||||
|
||||
if (request.stream) {
|
||||
// First let worker instantiate a generator
|
||||
const msg: WorkerRequest = {
|
||||
kind: "completionStreamInit",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
request: request,
|
||||
selectedModelId: selectedModelId,
|
||||
modelId: this.modelId,
|
||||
chatOpts: this.chatOpts,
|
||||
},
|
||||
};
|
||||
await this.getPromise<null>(msg);
|
||||
|
||||
// Then return an async chunk generator that resides on the client side
|
||||
return this.asyncGenerate(selectedModelId) as AsyncGenerator<
|
||||
Completion,
|
||||
void,
|
||||
void
|
||||
>;
|
||||
}
|
||||
|
||||
// Non streaming case is more straightforward
|
||||
const msg: WorkerRequest = {
|
||||
kind: "completionNonStreaming",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
request: request,
|
||||
modelId: this.modelId,
|
||||
chatOpts: this.chatOpts,
|
||||
},
|
||||
};
|
||||
return await this.getPromise<Completion>(msg);
|
||||
}
|
||||
|
||||
async embedding(
|
||||
request: EmbeddingCreateParams,
|
||||
): Promise<CreateEmbeddingResponse> {
|
||||
if (this.modelId === undefined) {
|
||||
throw new WorkerEngineModelNotLoadedError(this.constructor.name);
|
||||
}
|
||||
const msg: WorkerRequest = {
|
||||
kind: "embedding",
|
||||
uuid: crypto.randomUUID(),
|
||||
content: {
|
||||
request: request,
|
||||
modelId: this.modelId,
|
||||
chatOpts: this.chatOpts,
|
||||
},
|
||||
};
|
||||
return await this.getPromise<CreateEmbeddingResponse>(msg);
|
||||
}
|
||||
|
||||
onmessage(event: any) {
|
||||
let msg: WorkerResponse;
|
||||
if (event instanceof MessageEvent) {
|
||||
msg = event.data as WorkerResponse;
|
||||
} else {
|
||||
msg = event as WorkerResponse;
|
||||
}
|
||||
switch (msg.kind) {
|
||||
case "initProgressCallback": {
|
||||
if (this.initProgressCallback !== undefined) {
|
||||
this.initProgressCallback(msg.content as InitProgressReport);
|
||||
}
|
||||
return;
|
||||
}
|
||||
case "return": {
|
||||
const cb = this.pendingPromise.get(msg.uuid);
|
||||
if (cb === undefined) {
|
||||
throw Error("return from a unknown uuid msg=" + msg.uuid);
|
||||
}
|
||||
this.pendingPromise.delete(msg.uuid);
|
||||
cb(msg);
|
||||
return;
|
||||
}
|
||||
case "throw": {
|
||||
const cb = this.pendingPromise.get(msg.uuid);
|
||||
if (cb === undefined) {
|
||||
throw Error("return from a unknown uuid, msg=" + msg);
|
||||
}
|
||||
this.pendingPromise.delete(msg.uuid);
|
||||
cb(msg);
|
||||
return;
|
||||
}
|
||||
default: {
|
||||
const unknownMsg = msg as any;
|
||||
throw new UnknownMessageKindError(unknownMsg.kind, unknownMsg.content);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user