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,23 @@
+# WebLLM Logit Processor and Low-Level API Example
+
+This folder explains the usage of `LogitProcessor`, demonstrating how it can be used to
+manipulate the raw logits before sampling the token (e.g. setting certain tokens to `inf` or `-inf`).
+We demonstrate how to use it with and without a web worker, which can be toggled with `USE_WEB_WORKER`
+in `logit_processor.ts` (see `worker.ts` on how `LogitProcessor` plays a role there).
+
+We also demonstrate the usage of a low-level API `forwardTokenAndSample()`, which, unlike `chat.completions.create()`
+that assumes the usage is for autoregressive chatting, here we have more fine-grained control.
+
+See `my_logit_processor.ts` on how to customize your own logit processor. Here we make the logit
+of token 0 `100.0` manually, large enough that we should expect to always sample token 0, which
+is indeed the case if we observe the console log. We also demonstarte that a LogitProcessor can be
+stateful, and the state can also be cleaned with `LogitProcessor.resetState()`.
+
+To try it out, you can do the following steps under this folder
+
+```bash
+npm install
+npm start
+```
+
+Note if you would like to hack WebLLM core package, you can change web-llm dependencies as `"file:../.."`, and follow the build from source instruction in the project to build webllm locally. This option is only recommended if you would like to hack WebLLM core package.
@@ -0,0 +1,20 @@
+{
+  "name": "logit-processor",
+  "version": "0.1.0",
+  "private": true,
+  "scripts": {
+    "start": "parcel src/logit_processor.html  --port 8885",
+    "build": "parcel build src/logit_processor.html --dist-dir lib"
+  },
+  "devDependencies": {
+    "buffer": "^5.7.1",
+    "parcel": "^2.8.3",
+    "process": "^0.11.10",
+    "tslib": "^2.3.1",
+    "typescript": "^4.9.5",
+    "url": "^0.11.3"
+  },
+  "dependencies": {
+    "@mlc-ai/web-llm": "^0.2.84"
+  }
+}
@@ -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);
+};