Skip to content

Commit 8912a4d

Browse files
tqchentqchen
authored
Introduce builtin WebWorker support (#118)
This PR introduces webworker support to WebLLM package.
1 parent 6e0dc08 commit 8912a4d

18 files changed

+527
-23
lines changed

‎README.md

+44-2
Original file line numberDiff line numberDiff line change
@@ -65,8 +65,50 @@ async function main() {
6565
main();
6666
```
6767

68-
Finally, you can find a complete
69-
You can also find a complete chat app in [examples/simple-chat](examples/simple-chat/).
68+
### Using Web Worker
69+
70+
WebLLM comes with API support for WebWorker so you can hook
71+
the generation process into a separate worker thread so that
72+
the compute in the webworker won't disrupt the UI.
73+
74+
We first create a worker script that created a ChatModule and
75+
hook it up to a handler that handles requests.
76+
77+
```typescript
78+
// worker.ts
79+
import { ChatWorkerHandler, ChatModule } from "@mlc-ai/web-llm";
80+
81+
// Hookup a chat module to a worker handler
82+
const chat = new ChatModule();
83+
const handler = new ChatWorkerHandler(chat);
84+
self.onmessage = (msg: MessageEvent) => {
85+
handler.onmessage(msg);
86+
};
87+
```
88+
89+
Then in the main logic, we create a `ChatWorkerClient` that
90+
implements the same `ChatInterface`. The rest of the logic remains the same.
91+
92+
```typescript
93+
// main.ts
94+
import * as webllm from "@mlc-ai/web-llm";
95+
96+
async function main() {
97+
// Use a chat worker client instead of ChatModule here
98+
const chat = new webllm.ChatWorkerClient(new Worker(
99+
new URL('./worker.ts', import.meta.url),
100+
{type: 'module'}
101+
));
102+
// everything else remains the same
103+
}
104+
```
105+
106+
107+
### Build a ChatApp
108+
109+
You can find a complete
110+
a complete chat app example in [examples/simple-chat](examples/simple-chat/).
111+
70112

71113
## Customized Model Weights
72114

‎examples/README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -6,5 +6,5 @@ Please send a pull request if you find things that belongs to here.
66
## Tutorial Examples
77

88
- [get-started](get-started): minimum get started example.
9+
- [web-worker](web-worker): get started with web worker backed chat.
910
- [simple-chat](simple-chat): a mininum and complete chat app.
10-

‎examples/get-started/README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ To try it out, you can do the following steps
77
- `@mlc-ai/web-llm` points to a valid npm version e.g.
88
```js
99
"dependencies": {
10-
"@mlc-ai/web-llm": "^0.1.3"
10+
"@mlc-ai/web-llm": "^0.2.0"
1111
}
1212
```
1313
Try this option if you would like to use WebLLM without building it yourself.

‎examples/simple-chat/README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ chat app based on WebLLM. To try it out, you can do the following steps
77
- Option 1: `@mlc-ai/web-llm` points to a valid npm version e.g.
88
```js
99
"dependencies": {
10-
"@mlc-ai/web-llm": "^0.1.3"
10+
"@mlc-ai/web-llm": "^0.2.0"
1111
}
1212
```
1313
Try this option if you would like to use WebLLM.

‎examples/simple-chat/src/gh-config.js

+2-1
Original file line numberDiff line numberDiff line change
@@ -18,5 +18,6 @@ export default {
1818
"vicuna-v1-7b-q4f32_0": "https://raw.githubusercontent.com/mlc-ai/binary-mlc-llm-libs/main/vicuna-v1-7b-q4f32_0-webgpu.wasm",
1919
"RedPajama-INCITE-Chat-3B-v1-q4f32_0": "https://raw.githubusercontent.com/mlc-ai/binary-mlc-llm-libs/main/RedPajama-INCITE-Chat-3B-v1-q4f32_0-webgpu.wasm",
2020
"RedPajama-INCITE-Chat-3B-v1-q4f16_0": "https://raw.githubusercontent.com/mlc-ai/binary-mlc-llm-libs/main/RedPajama-INCITE-Chat-3B-v1-q4f16_0-webgpu.wasm"
21-
}
21+
},
22+
"use_web_worker": true
2223
}

‎examples/simple-chat/src/llm_chat.html

-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
<select id="chatui-select">
22
</select>
33

4-
54
<link href="./llm_chat.css" rel="stylesheet" type="text/css"/>
65

76
<div class="chatui">

‎examples/simple-chat/src/mlc-local-config.js

+2-1
Original file line numberDiff line numberDiff line change
@@ -22,5 +22,6 @@ export default {
2222
"vicuna-v1-7b-q4f32_0": "http://localhost:8000/vicuna-v1-7b-q4f32_0/vicuna-v1-7b-q4f32_0-webgpu.wasm",
2323
"RedPajama-INCITE-Chat-3B-v1-q4f32_0": "http://localhost:8000/RedPajama-INCITE-Chat-3B-v1-q4f32_0/RedPajama-INCITE-Chat-3B-v1-q4f32_0-webgpu.wasm",
2424
"RedPajama-INCITE-Chat-3B-v1-q4f16_0": "http://localhost:8000/RedPajama-INCITE-Chat-3B-v1-q4f16_0/RedPajama-INCITE-Chat-3B-v1-q4f16_0-webgpu.wasm"
25-
}
25+
},
26+
"use_web_worker": true
2627
}

‎examples/simple-chat/src/simple_chat.ts

+20-12
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import appConfig from "./app-config";
2-
import { ChatModule, ModelRecord } from "@mlc-ai/web-llm";
2+
import { ChatInterface, ChatModule, ChatWorkerClient, ModelRecord } from "@mlc-ai/web-llm";
33

44
function getElementAndCheck(id: string): HTMLElement {
55
const element = document.getElementById(id);
@@ -18,7 +18,7 @@ class ChatUI {
1818
private uiChat: HTMLElement;
1919
private uiChatInput: HTMLInputElement;
2020
private uiChatInfoLabel: HTMLLabelElement;
21-
private chat: ChatModule;
21+
private chat: ChatInterface;
2222
private config: AppConfig = appConfig;
2323
private selectedModel: string;
2424
private chatLoaded = false;
@@ -27,8 +27,9 @@ class ChatUI {
2727
// all requests send to chat are sequentialized
2828
private chatRequestChain: Promise<void> = Promise.resolve();
2929

30-
constructor() {
31-
this.chat = new ChatModule();
30+
constructor(chat: ChatInterface) {
31+
// use web worker to run chat generation in background
32+
this.chat = chat;
3233
// get the elements
3334
this.uiChat = getElementAndCheck("chatui-chat");
3435
this.uiChatInput = getElementAndCheck("chatui-input") as HTMLInputElement;
@@ -156,9 +157,10 @@ class ChatUI {
156157
private resetChatHistory() {
157158
const clearTags = ["left", "right", "init", "error"];
158159
for (const tag of clearTags) {
159-
const matches = this.uiChat.getElementsByClassName(`msg ${tag}-msg`);
160+
// need to unpack to list so the iterator don't get affected by mutation
161+
const matches = [...this.uiChat.getElementsByClassName(`msg ${tag}-msg`)];
160162
for (const item of matches) {
161-
item.remove();
163+
this.uiChat.removeChild(item);
162164
}
163165
}
164166
if (this.uiChatInfoLabel !== undefined) {
@@ -211,11 +213,6 @@ class ChatUI {
211213

212214
this.appendMessage("left", "");
213215
const callbackUpdateResponse = (step, msg) => {
214-
if (msg.endsWith("##")) {
215-
msg = msg.substring(0, msg.length - 2);
216-
} else if (msg.endsWith("#")) {
217-
msg = msg.substring(0, msg.length - 1);
218-
}
219216
this.updateLastMessage("left", msg);
220217
};
221218

@@ -233,4 +230,15 @@ class ChatUI {
233230
}
234231
}
235232

236-
new ChatUI();
233+
const useWebWorker = appConfig.use_web_worker;
234+
let chat: ChatInterface;
235+
236+
if (useWebWorker) {
237+
chat = new ChatWorkerClient(new Worker(
238+
new URL('./worker.ts', import.meta.url),
239+
{type: 'module'}
240+
));
241+
} else {
242+
chat = new ChatModule();
243+
}
244+
new ChatUI(chat);

‎examples/simple-chat/src/worker.ts

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
// Serve the chat workload through web worker
2+
import { ChatWorkerHandler, ChatModule } from "@mlc-ai/web-llm";
3+
4+
const chat = new ChatModule();
5+
const handler = new ChatWorkerHandler(chat);
6+
self.onmessage = (msg: MessageEvent) => {
7+
handler.onmessage(msg);
8+
};

‎examples/web-worker/README.md

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# WebLLM Get Started with WebWorker
2+
3+
This folder provides a minimum demo to show WebLLM API using
4+
[WebWorker](https://developer.mozilla.org/en-US/docs/Web/API/Web_Workers_API/Using_web_workers).
5+
The main benefit of web worker is that all ML workloads runs on a separate thread as a result
6+
will less likely block the UI.
7+
8+
To try it out, you can do the following steps
9+
10+
- Modify [package.json](package.json) to make sure either
11+
- `@mlc-ai/web-llm` points to a valid npm version e.g.
12+
```js
13+
"dependencies": {
14+
"@mlc-ai/web-llm": "^0.2.0"
15+
}
16+
```
17+
Try this option if you would like to use WebLLM without building it yourself.
18+
- Or keep the dependencies as `"file:../.."`, and follow the build from source
19+
instruction in the project to build webllm locally. This option is more useful
20+
for developers who would like to hack WebLLM core package.
21+
- Run the following command
22+
```bash
23+
npm install
24+
npm start
25+
```

‎examples/web-worker/package.json

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
{
2+
"name": "get-started-web-worker",
3+
"version": "0.1.0",
4+
"private": true,
5+
"scripts": {
6+
"start": "parcel src/get_started.html --port 8888",
7+
"build": "parcel build src/get_started.html --dist-dir lib"
8+
},
9+
"devDependencies": {
10+
"parcel": "^2.8.3",
11+
"typescript": "^4.9.5",
12+
"tslib": "^2.3.1"
13+
},
14+
"dependencies": {
15+
"@mlc-ai/web-llm": "file:../.."
16+
}
17+
}

‎examples/web-worker/src/get_started.html

+22
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
<!DOCTYPE html>
2+
<html>
3+
<script>
4+
webLLMGlobal = {}
5+
</script>
6+
<body>
7+
<h2>WebLLM Test Page</h2>
8+
Open console to see output
9+
</br>
10+
</br>
11+
<label id="init-label"> </label>
12+
13+
<h3>Prompt</h3>
14+
<label id="prompt-label"> </label>
15+
16+
<h3>Response</h3>
17+
<label id="generate-label"> </label>
18+
</br>
19+
<label id="stats-label"> </label>
20+
21+
<script type="module" src="./main.ts"></script>
22+
</html>

‎examples/web-worker/src/main.ts

+41
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import * as webllm from "@mlc-ai/web-llm";
2+
3+
function setLabel(id: string, text: string) {
4+
const label = document.getElementById(id);
5+
if (label == null) {
6+
throw Error("Cannot find label " + id);
7+
}
8+
label.innerText = text;
9+
}
10+
11+
async function main() {
12+
// Use a chat worker client instead of ChatModule here
13+
const chat = new webllm.ChatWorkerClient(new Worker(
14+
new URL('./worker.ts', import.meta.url),
15+
{type: 'module'}
16+
));
17+
18+
chat.setInitProgressCallback((report: webllm.InitProgressReport) => {
19+
setLabel("init-label", report.text);
20+
});
21+
22+
await chat.reload("vicuna-v1-7b-q4f32_0");
23+
24+
const generateProgressCallback = (_step: number, message: string) => {
25+
setLabel("generate-label", message);
26+
};
27+
28+
const prompt0 = "What is the capital of Canada?";
29+
setLabel("prompt-label", prompt0);
30+
const reply0 = await chat.generate(prompt0, generateProgressCallback);
31+
console.log(reply0);
32+
33+
const prompt1 = "Can you write a poem about it?";
34+
setLabel("prompt-label", prompt1);
35+
const reply1 = await chat.generate(prompt1, generateProgressCallback);
36+
console.log(reply1);
37+
38+
console.log(await chat.runtimeStatsText());
39+
}
40+
41+
main();

‎examples/web-worker/src/worker.ts

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import { ChatWorkerHandler, ChatModule } from "@mlc-ai/web-llm";
2+
3+
// Hookup a chat module to a worker handler
4+
const chat = new ChatModule();
5+
const handler = new ChatWorkerHandler(chat);
6+
self.onmessage = (msg: MessageEvent) => {
7+
handler.onmessage(msg);
8+
};

‎package-lock.json

+2-2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "@mlc-ai/web-llm",
3-
"version": "0.1.3",
3+
"version": "0.2.0",
44
"description": "Hardware accelerated language model chats on browsers",
55
"main": "lib/index.js",
66
"types": "lib/index.d.ts",

‎src/index.ts

+5
Original file line numberDiff line numberDiff line change
@@ -13,3 +13,8 @@ export {
1313
export {
1414
ChatModule,
1515
} from "./chat_module";
16+
17+
export {
18+
ChatWorkerHandler,
19+
ChatWorkerClient
20+
} from "./web_worker";

0 commit comments

Comments
 (0)