Spaces:
Running
Running
| // Browser side of the wllama benchmark. Runs inside the page that run.mjs | |
| // opens, so a model is loaded and prompted through MiniSearch's own modules: | |
| // the catalog entry, the warmup-then-load in initializeWllama, and the chat | |
| // messages the search flow builds from results and page excerpts. | |
| import { | |
| getSettings, | |
| queryPubSub, | |
| settingsPubSub, | |
| updateLlmTextSearchResults, | |
| updatePageContents, | |
| } from "../../client/modules/pubSub"; | |
| import { | |
| defaultContextSize, | |
| getDefaultChatMessages, | |
| getFormattedSearchResults, | |
| } from "../../client/modules/textGenerationUtilities"; | |
| import type { TextSearchResults } from "../../client/modules/types"; | |
| import { initializeWllama, wllamaModels } from "../../client/modules/wllama"; | |
| interface BenchCase { | |
| name: string; | |
| query: string; | |
| results: TextSearchResults; | |
| pages: Record<string, string>; | |
| } | |
| interface BenchOptions { | |
| /** A catalog id, or `owner/repo/file.gguf` for a model not in the catalog. */ | |
| model: string; | |
| device: "webgpu" | "cpu"; | |
| /** Defaults to MiniSearch's own `cpuThreads` setting. */ | |
| threads?: number; | |
| cases: BenchCase[]; | |
| chatTemplateKwargs?: Record<string, unknown>; | |
| reasoningBudgetTokens?: number; | |
| /** Seconds before an answer is cut off and marked as timed out. */ | |
| caseTimeoutSeconds: number; | |
| /** Replaces the default system prompt, to measure a change to it. */ | |
| systemPrompt?: string; | |
| } | |
| function resolveModel(id: string) { | |
| const catalogModel = wllamaModels[id]; | |
| if (catalogModel) return catalogModel; | |
| const [owner, repo, ...file] = id.split("/"); | |
| if (!owner || !repo || file.length === 0) { | |
| throw new Error( | |
| `"${id}" is neither a catalog id nor an owner/repo/file.gguf path`, | |
| ); | |
| } | |
| // Starts from a catalog entry's shared defaults, minus the thinking budget | |
| // that only some entries set, as a new catalog entry would. | |
| return { | |
| ...Object.values(wllamaModels)[0], | |
| reasoningBudgetTokens: undefined, | |
| label: id, | |
| hfRepoId: `${owner}/${repo}`, | |
| hfFilePath: file.join("/"), | |
| }; | |
| } | |
| async function runBench(options: BenchOptions) { | |
| if (options.device === "webgpu" && !(await navigator.gpu?.requestAdapter())) { | |
| throw new Error("This browser exposes no WebGPU adapter"); | |
| } | |
| const catalogModel = resolveModel(options.model); | |
| const model = { | |
| ...catalogModel, | |
| chatTemplateKwargs: | |
| options.chatTemplateKwargs ?? catalogModel.chatTemplateKwargs, | |
| reasoningBudgetTokens: | |
| options.reasoningBudgetTokens ?? catalogModel.reasoningBudgetTokens, | |
| }; | |
| if (options.systemPrompt) { | |
| settingsPubSub[0]({ ...getSettings(), systemPrompt: options.systemPrompt }); | |
| } | |
| const loadStart = performance.now(); | |
| // Mirrors initializeWllamaInstance in client/modules/textGenerationWithWllama.ts, | |
| // with the device forced instead of detected. | |
| const wllama = await initializeWllama(model.hfRepoId, model.hfFilePath, { | |
| wllama: { suppressNativeLog: false, allowOffline: true }, | |
| chatTemplateKwargs: model.chatTemplateKwargs, | |
| model: { | |
| n_threads: options.threads ?? getSettings().cpuThreads, | |
| n_ctx: model.contextSize, | |
| n_batch: 512, | |
| cache_type_k: model.cacheTypeK, | |
| cache_type_v: model.cacheTypeV, | |
| flash_attn: model.flash_attn, | |
| embeddings: false, | |
| n_gpu_layers: options.device === "webgpu" ? 99999 : 0, | |
| ...(model.reasoningBudgetTokens | |
| ? { reasoning_budget_tokens: model.reasoningBudgetTokens } | |
| : {}), | |
| }, | |
| }); | |
| const loadSeconds = (performance.now() - loadStart) / 1000; | |
| const cases = []; | |
| for (const benchCase of options.cases) { | |
| queryPubSub[0](benchCase.query); | |
| updateLlmTextSearchResults(benchCase.results); | |
| updatePageContents(benchCase.pages); | |
| const messages = getDefaultChatMessages( | |
| await getFormattedSearchResults(model.shouldIncludeUrlsOnPrompt), | |
| ); | |
| const start = performance.now(); | |
| let firstAnswerAt: number | undefined; | |
| let thinkingTokens = 0; | |
| let answerTokens = 0; | |
| let reasoning = ""; | |
| let answer = ""; | |
| let timings: Record<string, number> | undefined; | |
| const abortController = new AbortController(); | |
| const timeout = setTimeout( | |
| () => abortController.abort(), | |
| options.caseTimeoutSeconds * 1000, | |
| ); | |
| const stream = await wllama.createChatCompletion({ | |
| messages, | |
| max_tokens: defaultContextSize, | |
| stream: true, | |
| abortSignal: abortController.signal, | |
| ...model.getSampling(), | |
| ...(model.chatTemplateKwargs | |
| ? { chat_template_kwargs: model.chatTemplateKwargs } | |
| : {}), | |
| }); | |
| try { | |
| for await (const chunk of stream) { | |
| const delta = chunk.choices[0]?.delta as | |
| | { content?: string; reasoning_content?: string } | |
| | undefined; | |
| if (delta?.reasoning_content) { | |
| thinkingTokens++; | |
| reasoning += delta.reasoning_content; | |
| } | |
| if (delta?.content) { | |
| firstAnswerAt ??= performance.now(); | |
| answerTokens++; | |
| answer += delta.content; | |
| } | |
| // The engine's own timings ride on the stream's chunks; the last one | |
| // covers the whole completion. | |
| const chunkTimings = (chunk as { timings?: Record<string, number> }) | |
| .timings; | |
| if (chunkTimings) timings = chunkTimings; | |
| } | |
| } catch (error) { | |
| if (!abortController.signal.aborted) throw error; | |
| } finally { | |
| clearTimeout(timeout); | |
| } | |
| const end = performance.now(); | |
| cases.push({ | |
| name: benchCase.name, | |
| promptTokens: timings?.prompt_n, | |
| promptTokensPerSecond: timings?.prompt_per_second, | |
| generatedTokensPerSecond: timings?.predicted_per_second, | |
| thinkingTokens, | |
| answerTokens, | |
| secondsToFirstAnswerToken: | |
| firstAnswerAt === undefined | |
| ? undefined | |
| : (firstAnswerAt - start) / 1000, | |
| totalSeconds: (end - start) / 1000, | |
| timedOut: abortController.signal.aborted, | |
| // Kept so the scorer can tell an echoed instruction from an answer. | |
| prompt: messages[0].content, | |
| reasoning, | |
| answer, | |
| }); | |
| } | |
| await wllama.exit(); | |
| return { model: model.label, device: options.device, loadSeconds, cases }; | |
| } | |
| Object.assign(window, { runBench, benchReady: true }); | |