nanochat-webgpu

Running

File size: 22,517 Bytes

91c9bfd

<!DOCTYPE html>
<html lang="en">

<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>NanoChat WebGPU</title>
  <link rel="icon"
    href="data:image/svg+xml,<svg xmlns=%22http://www.w3.org/2000/svg%22 viewBox=%220 0 100 100%22><text y=%22.9em%22 font-size=%2290%22>⚡</text></svg>">
  <style>
    :root {
      color-scheme: light;
    }

    * {
      box-sizing: border-box;
    }

    body {
      font-family: ui-sans-serif, -apple-system, system-ui, "Segoe UI", Helvetica, "Apple Color Emoji", Arial, sans-serif, "Segoe UI Emoji", "Segoe UI Symbol";
      background-color: #ffffff;
      color: #111827;
      min-height: 100vh;
      margin: 0;
      display: flex;
      flex-direction: column;
      max-height: 100vh;
    }

    /* Loader Styles */
    .loader-overlay {
      position: fixed;
      inset: 0;
      background-color: rgba(255, 255, 255, 0.8);
      backdrop-filter: blur(8px);
      display: flex;
      flex-direction: column;
      align-items: center;
      justify-content: center;
      z-index: 1000;
      color: #374151;
    }

    .loader {
      border: 5px solid #f3f3f3;
      border-top: 5px solid #3498db;
      border-radius: 50%;
      width: 50px;
      height: 50px;
      animation: spin 1s linear infinite;
    }

    #loader-text {
      margin: 1rem 0 0.5rem 0;
      font-weight: 500;
    }

    @keyframes spin {
      0% {
        transform: rotate(0deg);
      }

      100% {
        transform: rotate(360deg);
      }
    }

    .header {
      background-color: #ffffff;
      padding: 1.25rem 1.5rem;
    }

    .header-left {
      display: flex;
      align-items: center;
      gap: 0.75rem;
    }

    .header-logo {
      height: 32px;
      width: auto;
    }

    .header h1 {
      font-size: 1.25rem;
      font-weight: 600;
      margin: 0;
      color: #111827;
    }

    .new-conversation-btn {
      width: 32px;
      height: 32px;
      padding: 0;
      border: 1px solid #e5e7eb;
      border-radius: 0.5rem;
      background-color: #ffffff;
      color: #6b7280;
      cursor: pointer;
      display: flex;
      align-items: center;
      justify-content: center;
      transition: all 0.2s ease;
    }

    .new-conversation-btn:hover {
      background-color: #f3f4f6;
      border-color: #d1d5db;
      color: #374151;
    }

    .chat-container {
      flex: 1;
      overflow-y: auto;
      background-color: #ffffff;
    }

    .chat-wrapper {
      max-width: 48rem;
      margin: 0 auto;
      padding: 2rem 1.5rem 3rem;
      display: flex;
      flex-direction: column;
      gap: 0.75rem;
    }

    .message {
      display: flex;
      justify-content: flex-start;
      margin-bottom: 0.5rem;
      color: #0d0d0d;
    }

    .message.assistant {
      justify-content: flex-start;
    }

    .message.user {
      justify-content: flex-end;
    }

    .message-content {
      white-space: pre-wrap;
      line-height: 1.6;
      max-width: 100%;
    }

    .message.assistant .message-content {
      background: transparent;
      border: none;
      padding: 0.25rem 0;
      cursor: pointer;
      border-radius: 0.5rem;
      padding: 0.5rem;
      margin-left: -0.5rem;
      transition: background-color 0.2s ease;
    }

    .message.assistant .message-content:hover {
      background-color: #f9fafb;
    }

    .message.user .message-content {
      background-color: #f3f4f6;
      border-radius: 1.25rem;
      padding: 0.8rem 1rem;
      max-width: 65%;
      cursor: pointer;
      transition: background-color 0.2s ease;
    }

    .message.user .message-content:hover {
      background-color: #e5e7eb;
    }

    .message.console .message-content {
      font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', 'Consolas', 'Courier New', monospace;
      font-size: 0.875rem;
      background-color: #fafafa;
      padding: 0.75rem 1rem;
      color: #374151;
      max-width: 80%;
    }

    .input-container {
      background-color: #ffffff;
      padding: 1rem;
    }

    .input-wrapper {
      max-width: 48rem;
      margin: 0 auto;
      display: flex;
      gap: 0.75rem;
      align-items: flex-end;
    }

    .chat-input {
      flex: 1;
      padding: 0.8rem 1rem;
      border: 1px solid #d1d5db;
      border-radius: 0.75rem;
      background-color: #ffffff;
      color: #111827;
      font-size: 1rem;
      line-height: 1.5;
      resize: none;
      outline: none;
      min-height: 54px;
      max-height: 200px;
      transition: border-color 0.2s ease, box-shadow 0.2s ease;
    }

    .chat-input::placeholder {
      color: #9ca3af;
    }

    .chat-input:focus {
      border-color: #2563eb;
      box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.1);
    }

    .send-button {
      flex-shrink: 0;
      padding: 0;
      width: 54px;
      height: 54px;
      border: 1px solid #111827;
      border-radius: 0.75rem;
      background-color: #111827;
      color: #ffffff;
      display: flex;
      align-items: center;
      justify-content: center;
      cursor: pointer;
      transition: background-color 0.2s ease, border-color 0.2s ease, color 0.2s ease;
    }

    .send-button:hover:not(:disabled) {
      background-color: #2563eb;
      border-color: #2563eb;
    }

    .send-button:disabled {
      cursor: not-allowed;
      border-color: #d1d5db;
      background-color: #e5e7eb;
      color: #9ca3af;
    }

    .error-message {
      background-color: #fee2e2;
      border: 1px solid #fecaca;
      color: #b91c1c;
      padding: 0.75rem 1rem;
      border-radius: 0.75rem;
      margin-top: 0.5rem;
    }

    .progress {
      font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', 'Consolas', 'Courier New', monospace;
      text-align: center;
    }
  </style>
</head>

<body>
  <div id="loader" class="loader-overlay">
    <div class="loader"></div>
    <p id="loader-text">Loading model...</p>
    <span class="progress">(0.00%)</span>
  </div>

  <div class="header">
    <div class="header-left">
      <button id="newConversationBtn" class="new-conversation-btn" title="New Conversation (Ctrl+Shift+N)">
        <svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"
          stroke-linecap="round" stroke-linejoin="round">
          <path d="M12 5v14"></path>
          <path d="M5 12h14"></path>
        </svg>
      </button>
      <h1>nanochat webgpu</h1>
    </div>
  </div>

  <div class="chat-container" id="chatContainer">
    <div class="chat-wrapper" id="chatWrapper">
    </div>
  </div>

  <div class="input-container">
    <div class="input-wrapper">
      <textarea id="chatInput" class="chat-input" placeholder="Ask anything" rows="1" disabled></textarea>
      <button id="sendButton" class="send-button" disabled>
        <svg data-icon="send" width="22" height="22" viewBox="0 0 24 24" fill="none" stroke="currentColor"
          stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
          <path d="M22 2L11 13"></path>
          <path d="M22 2l-7 20-4-9-9-4 20-7z"></path>
        </svg>
        <svg data-icon="stop" width="22" height="22" viewBox="0 0 24 24" fill="currentColor" stroke="currentColor"
          stroke-width="2" stroke-linecap="round" stroke-linejoin="round" style="display: none;">
          <rect x="6" y="6" width="12" height="12"></rect>
        </svg>
      </button>
    </div>
  </div>

  <script type="module">
    import { pipeline, TextStreamer, InterruptableStoppingCriteria } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.6';

    /**
     * Encapsulates the entire chat application logic.
     */
    class ChatApp {
      constructor() {
        // Application State
        this.generator = null;
        this.messages = [];
        this.isGenerating = false;
        this.stoppingCriteria = new InterruptableStoppingCriteria();

        // Generation Settings
        this.settings = {
          temperature: 0.05,
          topK: 5,
        };
        this.fileProgress = new Map();

        // UI Element References
        this.ui = {};

        this.cacheDOMElements();
        this.bindEventListeners();
        this.initialize();
      }

      /**
       * Caches all required DOM elements for easy access.
       */
      cacheDOMElements() {
        this.ui.chatContainer = document.getElementById('chatContainer');
        this.ui.chatWrapper = document.getElementById('chatWrapper');
        this.ui.chatInput = document.getElementById('chatInput');
        this.ui.sendButton = document.getElementById('sendButton');
        this.ui.newConversationBtn = document.getElementById('newConversationBtn');
        this.ui.loader = document.getElementById('loader');
        this.ui.loaderText = document.getElementById('loader-text');
        this.ui.progressSpan = this.ui.loader.querySelector('.progress');
        this.ui.sendIcon = this.ui.sendButton.querySelector('[data-icon="send"]');
        this.ui.stopIcon = this.ui.sendButton.querySelector('[data-icon="stop"]');
      }

      /**
       * Binds all event listeners for user interactions.
       */
      bindEventListeners() {
        this.ui.sendButton.addEventListener('click', () => this.handlePrimaryButtonClick());
        this.ui.newConversationBtn.addEventListener('click', () => this.handleNewConversation());

        this.ui.chatInput.addEventListener('keydown', (event) => {
          if (event.key === 'Enter' && !event.shiftKey) {
            event.preventDefault();
            this.handlePrimaryButtonClick();
          }
        });

        this.ui.chatInput.addEventListener('input', () => {
          // Auto-resize textarea
          this.ui.chatInput.style.height = 'auto';
          this.ui.chatInput.style.height = Math.min(this.ui.chatInput.scrollHeight, 200) + 'px';
          this.updateSendButtonState();
        });

        document.addEventListener('keydown', (event) => {
          if (event.ctrlKey && event.shiftKey && event.key === 'N') {
            event.preventDefault();
            if (!this.isGenerating) {
              this.handleNewConversation();
            }
          }
        });
      }

      /**
       * Starts the application by setting the initial state and loading the model.
       */
      initialize() {
        this.messages.push({ role: 'system', content: 'You are a helpful assistant.' });
        this.initializeModel();
      }

      /**
       * Loads the text-generation model using the transformers.js pipeline.
       */
      async initializeModel() {
        try {
          this.fileProgress.clear();
          this.generator = await pipeline(
            "text-generation",
            "onnx-community/nanochat-d32-ONNX",
            {
              dtype: "q4",
              device: "webgpu",
              progress_callback: (data) => {
                if (data.status !== 'progress') return;
                this.fileProgress.set(data.file, { loaded: data.loaded, total: data.total });
                if (this.fileProgress.size >= 7) {
                  const aggregate = [...this.fileProgress.values()].reduce((acc, { loaded, total }) => {
                    acc.loaded += loaded;
                    acc.total += total;
                    return acc;
                  }, { loaded: 0, total: 0 });
                  const percent = aggregate.total > 0 ? ((aggregate.loaded / aggregate.total) * 100).toFixed(2) : '0.00';
                  this.ui.progressSpan.textContent = `(${percent}%)`;
                }
              }
            }
          );
          this.ui.loader.style.display = 'none';
          this.ui.chatInput.disabled = false;
          this.ui.chatInput.focus();
        } catch (error) {
          console.error("Error initializing model:", error);
          this.ui.loaderText.textContent = `Error loading model: ${error.message}.`;
          const loaderSpinner = this.ui.loader.querySelector('.loader');
          if (loaderSpinner) loaderSpinner.style.display = 'none';
        }
      }

      /**
       * Handles the primary button click, which can be either "Send" or "Stop".
       */
      handlePrimaryButtonClick() {
        if (this.isGenerating) {
          this.handleStopGeneration();
        } else {
          this.handleSendMessage();
        }
      }

      /**
       * Interrupts the current text generation.
       */
      handleStopGeneration() {
        this.stoppingCriteria.interrupt();
      }

      /**
       * Handles sending a message, including slash commands.
       */
      handleSendMessage() {
        const message = this.ui.chatInput.value.trim();
        if (!message || this.isGenerating) return;

        if (message.startsWith('/')) {
          if (!this.handleSlashCommand(message)) {
            this.renderMessage('console', `Unknown command: ${message}`);
          }
        } else {
          const userMessageIndex = this.messages.length;
          this.messages.push({ role: 'user', content: message });
          this.renderMessage('user', message, userMessageIndex);
          this.generateAssistantResponse();
        }

        this.ui.chatInput.value = '';
        this.ui.chatInput.style.height = 'auto';
        this.updateSendButtonState();
      }

      /**
       * Handles the regeneration of an assistant's response.
       * @param {number} messageIndex - The index of the message to regenerate.
       */
      handleRegenerate(messageIndex) {
        if (this.isGenerating || messageIndex < 1 || messageIndex >= this.messages.length) return;
        const messageToRegenerate = this.messages[messageIndex];
        if (messageToRegenerate.role !== 'assistant') return;

        this.messages = this.messages.slice(0, messageIndex);

        const allMessages = this.ui.chatWrapper.querySelectorAll('.message');
        // DOM has one less message than the array (no system prompt)
        for (let i = messageIndex - 1; i < allMessages.length; i++) {
          allMessages[i].remove();
        }

        this.generateAssistantResponse();
      }

      /**
       * Allows editing a user message and continuing the conversation from there.
       * @param {number} messageIndex - The index of the user message to edit.
       */
      handleEdit(messageIndex) {
        if (this.isGenerating || messageIndex < 1 || messageIndex >= this.messages.length) return;
        const messageToEdit = this.messages[messageIndex];
        if (messageToEdit.role !== 'user') return;

        this.ui.chatInput.value = messageToEdit.content;
        this.ui.chatInput.style.height = 'auto';
        this.ui.chatInput.style.height = Math.min(this.ui.chatInput.scrollHeight, 200) + 'px';

        this.messages = this.messages.slice(0, messageIndex);

        const allMessages = this.ui.chatWrapper.querySelectorAll('.message');
        for (let i = messageIndex - 1; i < allMessages.length; i++) {
          allMessages[i].remove();
        }

        this.updateSendButtonState();
        this.ui.chatInput.focus();
      }

      /**
       * Clears the conversation and resets the state.
       */
      handleNewConversation() {
        if (this.isGenerating) return;
        this.messages = [{ role: 'system', content: 'You are a helpful assistant.' }];
        this.ui.chatWrapper.innerHTML = '';
        this.ui.chatInput.value = '';
        this.ui.chatInput.style.height = 'auto';
        this.updateSendButtonState();
        this.ui.chatInput.focus();
      }

      /**
       * Parses and executes slash commands.
       * @param {string} command - The full command string from the user.
       * @returns {boolean} - True if the command was valid, false otherwise.
       */
      handleSlashCommand(command) {
        const parts = command.trim().split(/\s+/);
        const cmd = parts[0].toLowerCase();
        const arg = parts[1];

        switch (cmd) {
          case '/temperature':
            if (arg === undefined) {
              this.renderMessage('console', `Current temperature: ${this.settings.temperature}`);
            } else {
              const temp = parseFloat(arg);
              if (isNaN(temp) || temp < 0 || temp > 2) {
                this.renderMessage('console', 'Invalid temperature. Must be between 0.0 and 2.0');
              } else {
                this.settings.temperature = temp;
                this.renderMessage('console', `Temperature set to ${this.settings.temperature}`);
              }
            }
            return true;
          case '/topk':
            if (arg === undefined) {
              this.renderMessage('console', `Current top-k: ${this.settings.topK}`);
            } else {
              const topk = parseInt(arg);
              if (isNaN(topk) || topk < 1 || topk > 200) {
                this.renderMessage('console', 'Invalid top-k. Must be between 1 and 200');
              } else {
                this.settings.topK = topk;
                this.renderMessage('console', `Top-k set to ${this.settings.topK}`);
              }
            }
            return true;
          case '/clear':
            this.handleNewConversation();
            return true;
          case '/help':
            this.renderMessage('console',
              'Available commands:\n' +
              '/temperature - Show current temperature\n' +
              '/temperature <value> - Set temperature (0.0-2.0)\n' +
              '/topk - Show current top-k\n' +
              '/topk <value> - Set top-k (1-200)\n' +
              '/clear - Clear conversation\n' +
              '/help - Show this help message'
            );
            return true;
          default:
            return false;
        }
      }

      /**
       * Generates a response from the model and streams it to the UI.
       */
      async generateAssistantResponse() {
        this.stoppingCriteria.reset();
        this.isGenerating = true;
        this.setButtonToStop();

        const assistantContentDiv = this.renderMessage('assistant', '');
        let fullResponse = '';

        const streamer = new TextStreamer(this.generator.tokenizer, {
          skip_prompt: true,
          skip_special_tokens: true,
          callback_function: (token) => {
            fullResponse += token;
            assistantContentDiv.textContent = fullResponse;
            this.ui.chatContainer.scrollTop = this.ui.chatContainer.scrollHeight;
          }
        });

        try {
          await this.generator(this.messages, {
            max_new_tokens: 512,
            do_sample: this.settings.temperature > 0,
            temperature: this.settings.temperature,
            top_k: this.settings.topK,
            repetition_penalty: 1.2,
            streamer,
            stopping_criteria: this.stoppingCriteria,
          });

          const assistantMessageIndex = this.messages.length;
          this.messages.push({ role: 'assistant', content: fullResponse });

          // Add click handler for regeneration after the message is complete
          assistantContentDiv.setAttribute('data-message-index', assistantMessageIndex);
          assistantContentDiv.setAttribute('title', 'Click to regenerate this response');
          assistantContentDiv.addEventListener('click', () => this.handleRegenerate(assistantMessageIndex));
        } catch (error) {
          // Don't show an error if it was a user interruption
          if (!error.message.includes('interrupted')) {
            console.error('Error during generation:', error);
            assistantContentDiv.textContent = '';
            const errorDiv = document.createElement('div');
            errorDiv.className = 'error-message';
            errorDiv.textContent = `Error: ${error.message}`;
            assistantContentDiv.appendChild(errorDiv);
          } else {
            // If interrupted, save the partial response
            const assistantMessageIndex = this.messages.length;
            this.messages.push({ role: 'assistant', content: fullResponse });
            assistantContentDiv.setAttribute('data-message-index', assistantMessageIndex);
            assistantContentDiv.setAttribute('title', 'Click to regenerate this response');
            assistantContentDiv.addEventListener('click', () => this.handleRegenerate(assistantMessageIndex));
          }
        } finally {
          this.isGenerating = false;
          this.setButtonToSend();
          this.ui.chatInput.focus();
        }
      }

      /**
       * Renders a message in the chat window.
       * @param {string} role - 'user', 'assistant', or 'console'.
       * @param {string} content - The message text.
       * @param {number|null} messageIndex - The index for user/assistant messages for edit/regenerate.
       * @returns {HTMLElement} - The content div of the new message.
       */
      renderMessage(role, content, messageIndex = null) {
        const messageDiv = document.createElement('div');
        messageDiv.className = `message ${role}`;

        const contentDiv = document.createElement('div');
        contentDiv.className = 'message-content';
        contentDiv.textContent = content;

        if (role === 'user' && messageIndex !== null) {
          contentDiv.setAttribute('data-message-index', messageIndex);
          contentDiv.setAttribute('title', 'Click to edit and restart from here');
          contentDiv.addEventListener('click', () => this.handleEdit(messageIndex));
        }

        messageDiv.appendChild(contentDiv);
        this.ui.chatWrapper.appendChild(messageDiv);

        this.ui.chatContainer.scrollTop = this.ui.chatContainer.scrollHeight;
        return contentDiv;
      }

      /**
       * Updates the enabled/disabled state of the send button.
       */
      updateSendButtonState() {
        this.ui.sendButton.disabled = !this.isGenerating && !this.ui.chatInput.value.trim();
      }

      /**
       * Changes the primary button to its "Stop" state.
       */
      setButtonToStop() {
        this.ui.sendIcon.style.display = 'none';
        this.ui.stopIcon.style.display = '';
        this.ui.sendButton.disabled = false;
        this.ui.sendButton.title = 'Stop generating';
      }

      /**
       * Changes the primary button back to its "Send" state.
       */
      setButtonToSend() {
        this.ui.sendIcon.style.display = '';
        this.ui.stopIcon.style.display = 'none';
        this.ui.sendButton.title = '';
        this.updateSendButtonState();
      }
    }

    // Initialize the application once the DOM is ready.
    new ChatApp();

  </script>
</body>

</html>