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

File tree

18 files changed

+527
-23
lines changed

18 files changed

+527
-23
lines changed

‎README.md

Lines changed: 44 additions & 2 deletions
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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 2 additions & 1 deletion
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

Lines changed: 0 additions & 1 deletion
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

Lines changed: 2 additions & 1 deletion
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

Lines changed: 20 additions & 12 deletions
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

Lines changed: 8 additions & 0 deletions
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

Lines changed: 25 additions & 0 deletions
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+
```

0 commit comments

Comments
 (0)