WebGPU GPT Model Demo

WebGPU GPT Model Demo
WebGPU is supported in your browser!

  PS: Loading models is 5x slower on the web rather than running locally. Just clone the repo and open!


Load Shakespeare Model
Load GPT2 117M Model
Destroy Cache


  Special models (download required):

Load GPT2 377M Model
Load GPT2 777M Model
Load GPT2 1.5B Model



Number of tokens:


Top K:


Temperature:


Generate Text

WILL:

Ah, how dare you challenge me? Have you forgotten I built WebGPT?

  const webgpuSupportMessage = document.getElementById("webgpuSupportMessage");
  const loadModelButton = document.getElementsByClassName("loadModelButton");
  const setModelButtonDisabled = (bool) => {
    for (let i = 0; i < loadModelButton.length; i++) loadModelButton[i].disabled = bool;
  };
  const destroyCacheButton = document.getElementById("destroyCacheButton");

  const tokensInput = document.getElementById("tokensInput");
  const topKInput = document.getElementById("topKInput");
  const temperatureInput = document.getElementById("temperatureInput");
  const generateButton = document.getElementById("generateButton");
  const promptTextarea = document.getElementById("prompt");

  let GPTModel = null;

  // Check for WebGPU support
  if (!navigator.gpu) {
    webgpuSupportMessage.innerHTML =
      "WebGPU is not supported in your browser yet. Update Chrome to v113 or download <a href='https://www.google.com/chrome/canary/'>Chrome Canary</a>. Also available on <a href='https://www.microsoftedgeinsider.com/en-us/download/canary'>Edge Canary</a>.";
    console.error("WebGPU is not supported");
  } else {
    webgpuSupportMessage.innerHTML = "WebGPU is supported in your browser!";
    setModelButtonDisabled(false);
  }

  async function startGeneration() {
    setTextareaDisabled(true);

    const prompt = promptTextarea.value || " ";
    const numTokens = tokensInput.value;

    const topK = topKInput.value;
    const temperature = temperatureInput.value;
    const textStream = GPTModel.generate(prompt, numTokens, topK, temperature);

    for await (const text of textStream) {
      promptTextarea.value += text;
    }

    setTextareaDisabled(false);
  }

  async function loadModel(folder, tokenizer) {
    setModelButtonDisabled(true);

    GPTModel = new GPT(folder, tokenizer);
    await GPTModel.initialize();

    promptTextarea.value = GPTModel.defaultPrompt;
    topKInput.value = GPTModel.defaultTopK;
    tokensInput.value = GPTModel.defaultTokens;
    temperatureInput.value = GPTModel.defaultTemperature;

    setTextareaDisabled(false);
    tokensInput.disabled = false;
    topKInput.disabled = false;
    temperatureInput.disabled = false;

    destroyCacheButton.disabled = false;
  }

  function setTextareaDisabled(bool) {
    promptTextarea.disabled = bool;
    generateButton.disabled = bool;
  }

  async function continueGeneration() {
    setTextareaDisabled(true);

    const prompt = outputDiv.innerHTML || " ";
    const numTokens = tokensInput.value;

    outputDiv.innerHTML = prompt;

    const topK = topKInput.value;
    const temperature = temperatureI