chore: import upstream snapshot with attribution

2026-07-13 12:42:51 +08:00
commit f73e710e38
276 changed files with 36598 additions and 0 deletions
@@ -0,0 +1,16 @@
+<!doctype html>
+<html>
+  <script>
+    webLLMGlobal = {};
+  </script>
+
+  <body>
+    <h2>WebLLM Logit Processor Test Page</h2>
+    Open console to see the effect of your logit processor.
+    <br />
+    <br />
+    <label id="init-label"> </label>
+
+    <script type="module" src="./logit_processor.ts"></script>
+  </body>
+</html>
@@ -0,0 +1,77 @@
+import * as webllm from "@mlc-ai/web-llm";
+import { MyLogitProcessor } from "./my_logit_processor";
+
+const USE_WEB_WORKER = true; // Toggle this to use Logit Processor without a web worker
+const AUTOREGRESS_LIMIT = 32; // How many tokens to generate for this test
+
+function setLabel(id: string, text: string) {
+  const label = document.getElementById(id);
+  if (label == null) {
+    throw Error("Cannot find label " + id);
+  }
+  label.innerText = text;
+}
+
+async function main() {
+  const initProgressCallback = (report: webllm.InitProgressReport) => {
+    setLabel("init-label", report.text);
+  };
+  // Instantiate myLogitProcessor, registering in the logitProcessorRegistry
+  const myLogitProcessor = new MyLogitProcessor();
+  const logitProcessorRegistry = new Map<string, webllm.LogitProcessor>();
+  logitProcessorRegistry.set("phi-2-q4f32_1-MLC", myLogitProcessor);
+
+  let engine: webllm.MLCEngineInterface;
+
+  // Depending on whether we use a web worker, the code is slightly different
+  if (USE_WEB_WORKER) {
+    // see worker.ts on how LogitProcessor plays a role there
+    engine = await webllm.CreateWebWorkerMLCEngine(
+      new Worker(new URL("./worker.ts", import.meta.url), { type: "module" }),
+      "phi-2-q4f32_1-MLC",
+      { initProgressCallback: initProgressCallback },
+    );
+  } else {
+    engine = await webllm.CreateMLCEngine("phi-2-q4f32_1-MLC", {
+      initProgressCallback: initProgressCallback,
+      logitProcessorRegistry: logitProcessorRegistry,
+    });
+  }
+
+  // Below we demonstrate the usage of a low-level API `forwardTokensAndSample()`
+  const prompt: Array<number> = [42];
+  let nextToken = await engine.forwardTokensAndSample(
+    prompt,
+    /*isPrefill=*/ true,
+  );
+  console.log(nextToken);
+
+  let counter = prompt.length;
+  while (counter < AUTOREGRESS_LIMIT) {
+    counter += 1;
+    nextToken = await engine.forwardTokensAndSample(
+      [nextToken],
+      /*isPrefill=*/ false,
+    );
+    console.log(nextToken);
+  }
+
+  // By calling `engine.resetChat()`, we triggers MyLogitProcessor.resetState()
+  engine.resetChat();
+  counter = prompt.length;
+  nextToken = await engine.forwardTokensAndSample(prompt, /*isPrefill=*/ true);
+  console.log(nextToken);
+  while (counter < AUTOREGRESS_LIMIT) {
+    counter += 1;
+    nextToken = await engine.forwardTokensAndSample(
+      [nextToken],
+      /*isPrefill=*/ false,
+    );
+    console.log(nextToken);
+  }
+
+  // `forwardTokensAndSample()` is made compatible with registering runtime stats.
+  console.log(await engine.runtimeStatsText());
+}
+
+main();
@@ -0,0 +1,21 @@
+import * as webllm from "@mlc-ai/web-llm";
+
+// Define LogitProcessor
+export class MyLogitProcessor implements webllm.LogitProcessor {
+  private tokenSequence: Array<number> = [];
+
+  processLogits(logits: Float32Array): Float32Array {
+    logits[0] = 100.0; // should be enough so that we always sample token 0 below
+    return logits;
+  }
+
+  processSampledToken(token: number): void {
+    this.tokenSequence.push(token);
+    console.log("processSampledToken: " + this.tokenSequence.length);
+  }
+
+  resetState(): void {
+    this.tokenSequence = [];
+    console.log("resetState");
+  }
+}
@@ -0,0 +1,15 @@
+// Serve the chat workload through web worker
+import * as webllm from "@mlc-ai/web-llm";
+import { MyLogitProcessor } from "./my_logit_processor";
+
+console.log("Use web worker for logit processor");
+
+const myLogitProcessor = new MyLogitProcessor();
+const logitProcessorRegistry = new Map<string, webllm.LogitProcessor>();
+logitProcessorRegistry.set("phi-2-q4f32_1-MLC", myLogitProcessor);
+
+const handler = new webllm.WebWorkerMLCEngineHandler();
+handler.setLogitProcessorRegistry(logitProcessorRegistry);
+self.onmessage = (msg: MessageEvent) => {
+  handler.onmessage(msg);
+};