Add custom grammar support

index.html

@@ -33,6 +33,31 @@
             ></label>
           </form>
         </div>
+        <div class="card">
+          <form>
+            <label for="grammar-selection" class="container"
+              ><span><b>Grammar</b></span>
+              <select id="grammar-selection" value="json">
+                <option value="json">JSON (Default)</option>
+                <option value="custom">Custom EBNF Grammar</option>
+              </select></label
+            >
+          </form>
+        </div>
+        <div id="ebnf-grammar-container" class="card hidden">
+          <form>
+            <label for="ebnf-grammar" class="container"
+              ><span><b>Custom EBNF Grammar</b></span>
+              <textarea
+                id="ebnf-grammar"
+                dir="ltr"
+                placeholder="Type your custom EBNF grammar..."
+                rows="1"
+              ></textarea
+              >
+            </label>
+          </form>
+        </div>
         <div class="card">
           <form>
             <label class="container"

src/index.js

@@ -2,151 +2,188 @@ import { prebuiltAppConfig, CreateMLCEngine } from "@charliefruan/web-llm";
 import hljs from "highlight.js";
 import ace from "ace-builds";
 
-// required for ace to resolve module correctly
+// 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');
+const { Type } = require("@sinclair/typebox");
 
 let engine = null;
+let useCustomGrammar = false;
+let customGrammar = 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]*`;
 
-const availableModels = prebuiltAppConfig.model_list
-  .filter(
-    (m) =>
-      m.model_id.startsWith("Llama-3-8B") ||
-      m.model_id.startsWith("Llama-3-70B") ||
-      m.model_id.startsWith("Llama-3.1-8B") ||
-      m.model_id.startsWith("Llama-3.1-70B") ||
-      m.model_id.startsWith("Llama-3.2-3B") ||
-      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;
-  document.getElementById("model-selection").appendChild(option);
-});
-document.getElementById("model-selection").value = selectedModel;
-document.getElementById("model-selection").onchange = (e) => {
-  selectedModel = e.target.value;
-  engine = null;
-};
-document.getElementById(
-  "prompt"
-).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.
-`;
-
-// JSON editor setup
-const editor = ace.edit("schema", {
-  //   mode: "ace/mode/javascript",
-  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(),
-})`);
-
-// Generate button
-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);
-        document.getElementById("output").textContent = progress.text;
-      },
-    });
-  }
-  const request = {
-    stream: true,
-    stream_options: { include_usage: true },
-    messages: [
-      {
-        role: "user",
-        content: document.getElementById("prompt").value,
-      },
-    ],
-    max_tokens: 128,
-    response_format: {
-      type: "json_object",
-      schema: schema,
-    },
+document.addEventListener("DOMContentLoaded", () => {
+  // Ensure elements are loaded before using them
+  const grammarSelection = document.getElementById("grammar-selection");
+  const ebnfContainer = document.getElementById("ebnf-grammar-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 = customGrammar;
+
+  // Update grammar value on change
+  ebnfTextarea.onchange = (ev) => {
+    customGrammar = ev.target.value;
+    console.log("Grammar updated: ", customGrammar);
+  };
+
+  // Handle grammar selection changes
+  grammarSelection.onchange = (ev) => {
+    console.log("Grammar selection changed:", ev.target.value);
+    if (ev.target.value === "json") {
+      ebnfContainer.classList.add("hidden");
+      useCustomGrammar = false;
+    } else {
+      ebnfContainer.classList.remove("hidden");
+      useCustomGrammar = true;
+    }
+  };
+
+  // Populate model selection dropdown
+  const availableModels = prebuiltAppConfig.model_list
+    .filter(
+      (m) =>
+        m.model_id.startsWith("Llama-3-8B") ||
+        m.model_id.startsWith("Llama-3-70B") ||
+        m.model_id.startsWith("Llama-3.1-8B") ||
+        m.model_id.startsWith("Llama-3.1-70B") ||
+        m.model_id.startsWith("Llama-3.2-3B") ||
+        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
   };
 
-  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;
+  // 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.`;
+
+  // 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;
     }
-    if (chunk.usage) {
-      usage = chunk.usage;
+    const schema = JSON.stringify(T);
+
+    if (!engine) {
+      engine = await CreateMLCEngine(selectedModel, {
+        initProgressCallback: (progress) => {
+          console.log(progress);
+          outputDiv.textContent = progress.text;
+        },
+      });
     }
-    document.getElementById("output").textContent = curMessage;
-  }
-  const finalMessage = await engine.getMessage();
-  console.log(finalMessage);
-  if (hljs) {
-    document.getElementById("output").innerHTML = hljs.highlight(finalMessage, {
+
+    const request = {
+      stream: true,
+      stream_options: { include_usage: true },
+      messages: [{ role: "user", content: promptTextarea.value }],
+      max_tokens: 128,
+      response_format: useCustomGrammar
+        ? { type: "grammar", grammar: customGrammar }
+        : { 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;
-  } else {
-    document.getElementById("output").textContent = finalMessage;
-  }
-  if (usage) {
-    const stats = usage['extra'];
-    document.getElementById("stats").textContent = `Prefill: ${stats['prefill_tokens_per_s'].toFixed(2)}, Decode: ${stats['decode_tokens_per_s'].toFixed(2)}, Grammar Init: ${stats['grammar_init_ms'].toFixed(2)}, Grammar Overhead Per Token: ${stats['grammar_per_token_ms'].toFixed(2)}`;
-    document.getElementById("stats").classList.remove("hidden");
-  }
-};
+
+    if (usage) {
+      statsParagraph.textContent = `Prefill: ${usage.extra.prefill_tokens_per_s.toFixed(
+        2
+      )}, Decode: ${usage.extra.decode_tokens_per_s.toFixed(
+        2
+      )}, Grammar Init Overhead: ${usage.extra.grammar_init_ms.toFixed(
+        2
+      )}, Grammar Per Token Overhead: ${usage.extra.grammar_per_token_ms.toFixed(
+        2
+      )}`;
+      statsParagraph.classList.remove("hidden");
+    }
+  };
+});

style/style.css

@@ -180,6 +180,10 @@ button:hover {
   cursor: pointer;
 }
 
+#ebnf-grammar {
+  height: 24rem;
+}
+
 @media (min-width: 1600px) {
   body {
     padding: 4rem;