web-ai-demos/perf-worker-gemma/src/script.js at main · GoogleChromeLabs/web-ai-demos · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
const url = new URL('./worker.js', import.meta.url);
import { MESSAGE_CODE, MODEL_STATUS } from './consts.js';

function displayModelStatus(status) {
  document.getElementById('modelStatus').className = status;
}

function runLLMInference() {
  const userPrompt = document.getElementById('userPrompt').value;
  worker.postMessage(userPrompt);
}

displayModelStatus(MODEL_STATUS.NOT_STARTED);
const worker = new Worker(url);

worker.onmessage = function (message) {
  console.info('[Main thread] 📬 Message from worker: ', message);

  if (!message.data || !message.data.code) {
    throw new Error(
      `Message from worker is empty or doesn't contain a code field: ${message}`
    );
  }
  const messageCode = message.data.code;
  switch (messageCode) {
    case MESSAGE_CODE.PREPARING_MODEL:
      displayModelStatus(MODEL_STATUS.PREPARING);
      document.getElementById('inferenceButton').disabled = true;
      break;

    case MESSAGE_CODE.MODEL_READY:
      displayModelStatus(MODEL_STATUS.READY);
      document.getElementById('inferenceButton').disabled = false;
      break;

    case MESSAGE_CODE.GENERATING_RESPONSE:
      displayModelStatus(MODEL_STATUS.GENERATING);
      document.getElementById('inferenceButton').disabled = true;
      break;

    case MESSAGE_CODE.RESPONSE_READY:
      displayModelStatus(MODEL_STATUS.READY);
      document.getElementById('inferenceButton').disabled = false;
      document.getElementById('llmOutput').innerText = message.data.payload;
      break;

    case MESSAGE_CODE.MODEL_ERROR:
      displayModelStatus(MODEL_STATUS.ERROR);
      document.getElementById('inferenceButton').disabled = true;
      break;

    case MESSAGE_CODE.INFERENCE_ERROR:
      displayModelStatus(MODEL_STATUS.ERROR);
      document.getElementById('inferenceButton').disabled = false;
      break;

    default:
      throw new Error(
        `Message from worker contains an unknown message code: ${messageCode}`
      );
  }
};

window.runLLMInference = runLLMInference;