src/index.js

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");

 // Initialize the custom grammar textarea
 ebnfTextarea.value = 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]*`;

 // 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("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
 };

 // JSON editor setup with Ace
 const editor = ace.edit("schema", {
 mode: "ace/mode/javascript",
 theme: "ace/theme/github",
 wrap: true,
 });

 editor.setTheme("ace/theme/github");
 editor.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(),
 })`);

 // 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 () => {
 const schemaInput = editor.getValue();
 let T;
 try {
 T = eval(schemaInput);
 } catch (e) {
 console.error("Invalid schema", e);
 return;
 }
 const schema = JSON.stringify(T);

 if (!engine) {
 engine = await CreateMLCEngine(selectedModel, {
 initProgressCallback: (progress) => {
 console.log(progress);
 outputDiv.textContent = progress.text;
 },
 });
 }

 const request = {
 stream: true,
 stream_options: { include_usage: true },
 messages: [{ role: "user", content: promptTextarea.value }],
 max_tokens: 128,
 response_format: useCustomGrammar
 ? { type: "grammar", grammar: ebnfTextarea.value }
 : { type: "json_object", schema: schema },
 };

 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 = [];
 if (usage.extra.prefill_tokens_per_s) {
 statsTextParts.push(`Prefill: ${usage.extra.prefill_tokens_per_s.toFixed(
 1
 )} tok/s`);
 }
 if (usage.extra.decode_tokens_per_s) {
 statsTextParts.push(`Decode: ${usage.extra.decode_tokens_per_s.toFixed(
 1
 )} tok/s`);
 }
 if (usage.extra.grammar_init_ms) {
 statsTextParts.push(`Grammar Init Overhead: ${usage.extra.grammar_init_ms.toFixed(
 0
 )} ms`);
 }
 if (usage.extra.grammar_per_token_ms) {
 statsTextParts.push(`Grammar Per Token Overhead: ${usage.extra.grammar_per_token_ms.toFixed(
 2
 )} ms`);
 }
 statsParagraph.textContent = statsTextParts.join(", ");
 statsParagraph.classList.remove("hidden");
 }
 };
});