Spaces:

mlc-ai
/

WebLLM-Structured-Generation-Playground

Running

WebLLM-Structured-Generation-Playground / src /index.js

Update src/index.js

0274b5f verified almost 2 years ago

7.55 kB

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

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