src/index.js

import { prebuiltAppConfig, CreateMLCEngine } from "@mlc-ai/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("SmolLM2")
    )
    .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(`{
   "title":"User",
   "type":"object",
   "properties":{
      "first_name":{
         "type":"string"
      },
      "last_name":{
         "type":"string"
      },
      "age":{
         "type":"integer"
      },
      "is_active":{
         "type":"boolean"
      }
   },
   "required":[
      "first_name",
      "last_name",
      "age"
   ]
}
`);

  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 = `
  Create a user profile for a sales person with following properties:
  - first_name: string
  - last_name: string
  - age: integer
  - is_active: boolean
`;
  // 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(JSON.parse(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) {
        console.log(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");
    }
  };
});