index.html

<!DOCTYPE html>
<html>
  <head>
    <title>WebGPU GPT Model Demo</title>
    <script src="helpers.js"></script>
    <script src="tokenizer.js"></script>
    <script src="kernels.js"></script>
    <script src="main.js"></script>
  </head>
  <body>
    <h1>WebGPU GPT Model Demo</h1>
    <p id="webgpuSupportMessage">Checking WebGPU support...</p>
    <p>
      <i>PS: Loading models is 5x slower on the web rather than running locally. Just <a href="https://github.com/0hq/WebGPT">clone the repo</a> and open!</i>
    </p>

    <button id="loadModelButton" onclick="loadModelHandler()" disabled>Load Shakespeare Model</button>
    <button id="loadGPT2ModelButton" onclick="loadGPT2ModelHandler()" disabled>Load GPT2 Model</button><br />
    <br />
    <label for="tokens">Number of tokens:</label>
    <input type="number" min="1" max="300" value="100" id="tokensInput" disabled />
    <br /><br />
    <label for="topK">Top K:</label>
    <input type="number" min="1" max="100" value="2" id="topKInput" disabled />
    <br /><br />
    <label for="temperature">Temperature:</label>
    <input type="number" step="0.01" min="0.1" max="2" value="1" id="temperatureInput" disabled />
    <br /><br />
    <button id="generateButton" onclick="startGeneration()" disabled>Generate Text</button>
    <br /><br />
    <textarea id="prompt" rows="15" cols="50" disabled>
WILL:
Ah, how dare you challenge me?
Have you forgotten I built WebGPT?&#13;&#10;</textarea
    >
    <br /><br />
    <script>
      const webgpuSupportMessage = document.getElementById("webgpuSupportMessage");
      const loadModelButton = document.getElementById("loadModelButton");
      const loadGPT2ModelButton = document.getElementById("loadGPT2ModelButton");

      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 Chrome Canary here: <a href='https://www.google.com/chrome/canary/'>https://www.google.com/chrome/canary/</a>";
        console.error("WebGPU is not supported");
      } else {
        webgpuSupportMessage.innerHTML = "WebGPU is supported in your browser!";
        loadModelButton.disabled = false;
        loadGPT2ModelButton.disabled = 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 loadModelHandler() {
        loadModelButton.disabled = true;
        loadGPT2ModelButton.disabled = true;

        GPTModel = new GPT("better_shakespeare", "char");
        await GPTModel.initialize();

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

      async function loadGPT2ModelHandler() {
        loadModelButton.disabled = true;
        loadGPT2ModelButton.disabled = true;

        GPTModel = new GPT("gpt2", "bpe");
        await GPTModel.initialize();

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

        // Set the default prompt for GPT2 model
        promptTextarea.value = "What is the answer to life, the universe, and everything?\n";

        topKInput.value = 5;
        tokensInput.value = 5;
        temperatureInput.value = 1;
      }

      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 = temperatureInput.value;
        const textStream = generate(prompt, numTokens, 10, topK, temperature);

        for await (const text of textStream) {
          outputDiv.innerHTML += text;
        }

        setTextareaDisabled(false);
      }
    </script>
  </body>
</html>