File size: 7,554 Bytes
ab229ee
 
 
 
5254a31
ab229ee
 
 
 
 
5254a31
4879fd8
 
5254a31
4879fd8
5254a31
 
 
 
cc1f2fe
5254a31
 
 
 
 
 
 
 
 
 
 
cc1f2fe
5254a31
 
 
cc1f2fe
5254a31
 
 
 
 
 
 
 
0274b5f
 
34c1901
5254a31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4879fd8
 
f18fa7e
 
5254a31
 
 
 
f18fa7e
 
5254a31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f18fa7e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5254a31
34c1901
 
 
 
 
 
 
 
 
5254a31
 
 
 
 
 
 
 
 
19eaf6a
f18fa7e
 
 
 
 
 
 
 
 
 
 
 
 
 
5254a31
 
 
 
4b86d74
f18fa7e
5254a31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4879fd8
 
5254a31
 
6508682
4b86d74
a394f47
4b86d74
6508682
a394f47
 
 
4b86d74
6508682
a394f47
 
668e490
48dfd32
6508682
 
a394f47
668e490
48dfd32
668e490
 
 
 
 
 
 
 
 
 
a394f47
6508682
a394f47
 
5254a31
 
 
 
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
import { prebuiltAppConfig, CreateMLCEngine } from "@charliefruan/web-llm";
import hljs from "highlight.js";
import ace from "ace-builds";

// Required for ace to resolve the module correctly
require("ace-builds/src-noconflict/mode-javascript");
require("ace-builds/webpack-resolver");

// DO NOT REMOVE
// Required for user input type definition to be eval
const { Type } = require("@sinclair/typebox");

let engine = null;
let useCustomGrammar = false;

document.addEventListener("DOMContentLoaded", () => {
  // Ensure elements are loaded before using them
  const grammarSelection = document.getElementById("grammar-selection");
  const ebnfContainer = document.getElementById("ebnf-grammar-container");
  const schemaContainer = document.getElementById("schema-container");
  const modelSelection = document.getElementById("model-selection");
  const ebnfTextarea = document.getElementById("ebnf-grammar");
  const promptTextarea = document.getElementById("prompt");
  const outputDiv = document.getElementById("output");
  const statsParagraph = document.getElementById("stats");

  // Handle grammar selection changes
  grammarSelection.onchange = (ev) => {
    console.log("Grammar selection changed:", ev.target.value);
    if (ev.target.value === "json") {
      ebnfContainer.classList.add("hidden");
      schemaContainer.classList.remove("hidden");
      useCustomGrammar = false;
    } else {
      ebnfContainer.classList.remove("hidden");
      schemaContainer.classList.add("hidden");
      useCustomGrammar = true;
    }
  };

  // Populate model selection dropdown
  const availableModels = prebuiltAppConfig.model_list
    .filter(
      (m) =>
        m.model_id.startsWith("Qwen2.5-0.5B-Instruct") ||
        m.model_id.startsWith("Qwen2.5-1.5B-Instruct") ||
        m.model_id.startsWith("Llama-3") ||
        m.model_id.startsWith("Hermes-2") ||
        m.model_id.startsWith("Hermes-3") ||
        m.model_id.startsWith("Phi-3")
    )
    .map((m) => m.model_id);

  let selectedModel = availableModels[0];

  availableModels.forEach((modelId) => {
    const option = document.createElement("option");
    option.value = modelId;
    option.textContent = modelId;
    modelSelection.appendChild(option);
  });

  modelSelection.value = selectedModel;

  modelSelection.onchange = (e) => {
    selectedModel = e.target.value;
    engine = null; // Reset the engine when the model changes
  };

  // Editors setup with Ace
  const jsonSchemaEditor = ace.edit("schema", {
    mode: "ace/mode/javascript",
    theme: "ace/theme/github",
    wrap: true,
  });
  jsonSchemaEditor.setTheme("ace/theme/github");
  jsonSchemaEditor.setValue(`Type.Object({
    "name": Type.String(),
    "house": Type.Enum({
      "Gryffindor": "Gryffindor",
      "Hufflepuff": "Hufflepuff",
      "Ravenclaw": "Ravenclaw",
      "Slytherin": "Slytherin",
    }),
    "blood_status": Type.Enum({
      "Pure-blood": "Pure-blood",
      "Half-blood": "Half-blood",
      "Muggle-born": "Muggle-born",
    }),
    "occupation": Type.Enum({
      "Student": "Student",
      "Professor": "Professor",
      "Ministry of Magic": "Ministry of Magic",
      "Other": "Other",
    }),
    "wand": Type.Object({
      "wood": Type.String(),
      "core": Type.String(),
      "length": Type.Number(),
    }),
    "alive": Type.Boolean(),
    "patronus": Type.String(),
  })`);

  const grammarEditor = ace.edit("ebnf-grammar", {
    theme: "ace/theme/github",
    wrap: true,
  });
  grammarEditor.setTheme("ace/theme/github");
  grammarEditor.setValue(String.raw`main ::= basic_array | basic_object
basic_any ::= basic_number | basic_string | basic_boolean | basic_null | basic_array | basic_object
basic_integer ::= ("0" | "-"? [1-9] [0-9]*) ".0"?
basic_number ::= ("0" | "-"? [1-9] [0-9]*) ("." [0-9]+)? ([eE] [+-]? [0-9]+)?
basic_string ::= (([\"] basic_string_1 [\"]))
basic_string_1 ::= "" | [^"\\\x00-\x1F] basic_string_1 | "\\" escape basic_string_1
escape ::= ["\\/bfnrt] | "u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9]
basic_boolean ::= "true" | "false"
basic_null ::= "null"
basic_array ::= "[" ("" | ws basic_any (ws "," ws basic_any)*) ws "]"
basic_object ::= "{" ("" | ws basic_string ws ":" ws basic_any ( ws "," ws basic_string ws ":" ws basic_any)*) ws "}"
ws ::= [\n\t]*`);

  // Set initial prompt
  promptTextarea.value = `Hermione Granger is a character in Harry Potter. Please fill in the following information about this character in JSON format.
Name is a string of character name.
House is one of Gryffindor, Hufflepuff, Ravenclaw, Slytherin.
Blood status is one of Pure-blood, Half-blood, Muggle-born.
Occupation is one of Student, Professor, Ministry of Magic, Other.
Wand is an object with wood, core, and length.
Alive is a boolean.
Patronus is a string.
`;
  // Generate button click handler
  document.getElementById("generate").onclick = async () => {
    if (!engine) {
      engine = await CreateMLCEngine(selectedModel, {
        initProgressCallback: (progress) => {
          console.log(progress);
          outputDiv.textContent = progress.text;
        },
      });
    }
    let response_format = { type: "grammar", grammar: grammarEditor.getValue() };
    if (!useCustomGrammar) {
      const schemaInput = jsonSchemaEditor.getValue();
      let T;
      try {
        T = eval(schemaInput);
      } catch (e) {
        console.error("Invalid schema", e);
        return;
      }
      const schema = JSON.stringify(T);
      response_format = { type: "json_object", schema }
    }
    console.log(response_format);
    const request = {
      stream: true,
      stream_options: { include_usage: true },
      messages: [{ role: "user", content: promptTextarea.value }],
      max_tokens: 512,
      response_format,
    };

    let curMessage = "";
    let usage = null;
    const generator = await engine.chatCompletion(request);

    for await (const chunk of generator) {
      const curDelta = chunk.choices[0]?.delta.content;
      if (curDelta) curMessage += curDelta;
      if (chunk.usage) usage = chunk.usage;
      outputDiv.textContent = curMessage;
    }

    const finalMessage = await engine.getMessage();
    outputDiv.innerHTML = hljs.highlight(finalMessage, {
      language: "json",
    }).value;

    if (usage) {
      const statsTextParts = [];
      console.log(usage);
      if (usage.extra.prefill_tokens_per_s) {
        statsTextParts.push(`Prefill Speed: ${usage.extra.prefill_tokens_per_s.toFixed(
          1
        )} tok/s`);
      }
      if (usage.extra.decode_tokens_per_s) {
        statsTextParts.push(`Decode Speed: ${usage.extra.decode_tokens_per_s.toFixed(
          1
        )} tok/s`);
      }
      if (usage.extra.time_per_output_token_s) {
        statsTextParts.push(`Time Per Output Token: ${(1000 * usage.extra.time_per_output_token_s).toFixed(
          0
        )} ms`);
      }
      if (usage.extra.time_to_first_token_s) {
        statsTextParts.push(`Time to First Token: ${(1000 * usage.extra.time_to_first_token_s).toFixed(
          0
        )} ms`);
      }
      if (usage.extra.grammar_init_s) {
        statsTextParts.push(`Grammar Init Overhead: ${(1000 * usage.extra.grammar_init_s).toFixed(
          0
        )} ms`);
      }
      if (usage.extra.grammar_per_token_s) {
        statsTextParts.push(`Grammar Per-token Overhead: ${(1000 * usage.extra.grammar_per_token_s).toFixed(
          2
        )} ms`);
      }
      statsParagraph.textContent = statsTextParts.join(", ");
      statsParagraph.classList.remove("hidden");
    }
  };
});