Chatroom

In this blog post, we'll explore how to build a web-based chat application that leverages small Large Language Models (LLMs) running directly in the browser using the mlc-llm package and WebGPU. We'll walk through a React component that demonstrates how to integrate these technologies to create a responsive and interactive chat interface.

Introduction

The advent of WebGPU has opened up new possibilities for running machine learning models directly in the browser, enabling applications that were previously only possible on the server side. The mlc-llm package allows developers to run quantized LLMs efficiently in the browser, providing a way to build privacy-preserving and low-latency applications.

In this tutorial, we'll dissect a React component that:

Loads small LLMs dynamically in the browser.
Provides a chat interface where users can interact with the model.
Utilizes WebGPU for efficient on-device computation.

Let's dive into the code and understand how these components come together.

Prerequisites

Basic understanding of React and TypeScript.
Familiarity with asynchronous programming in JavaScript.
Knowledge of WebGPU and LLMs is helpful but not required.

The Code Overview

The core of our application is the Chat component, which manages the user interface and the interaction with the LLM. Here's a breakdown of the main parts:

State Management: Using React hooks to manage the state of the application.
Model Loading: Dynamically loading LLMs with mlc-llm and integrating them with WebGPU.
Chat Interface: Handling user input and displaying messages.
Progress Monitoring: Providing feedback during model loading.

Importing Dependencies

"use client";

import { useEffect, useRef, useState } from "react";
import {
  ChatCompletionAssistantMessageParam,
  ChatCompletionChunk,
  ChatCompletionMessageParam,
  ChatCompletionUserMessageParam,
  CreateMLCEngine,
  MLCEngine,
} from "@mlc-ai/web-llm";
import ReactMarkdown from "react-markdown";

React Hooks: We use useState, useEffect, and useRef for state and lifecycle management.
MLC-LLM: Importing necessary types and functions from the @mlc-ai/web-llm package.
Markdown Rendering: ReactMarkdown is used to render messages that may contain Markdown syntax.

Defining Types and Utility Functions

type ChatMessage = {
  role: "assistant" | "user";
  message: string;
};

type Progress = {
  progress: number;
  timeElapsed: number;
  text: string;
};

const mapChatMessageToChatCompletionParam = (
  messages: ChatMessage[]
): Array<
  ChatCompletionUserMessageParam | ChatCompletionAssistantMessageParam
> =>
  messages.map((message) => ({
    role: message.role,
    content: message.message,
  }));

ChatMessage: Represents a message in the chat, either from the user or the assistant.
Progress: Keeps track of the loading or processing progress.
mapChatMessageToChatCompletionParam: Converts our ChatMessage objects to the format expected by the mlc-llm API.

The Chat Component

State Variables

const [messages, setMessages] = useState<ChatMessage[]>([]);
const [input, setInput] = useState("");
const [model, setModel] = useState("Qwen2.5-3B-Instruct-q4f32_1-MLC");
const [loadedModel, setLoadedModel] = useState<null | string>(null);
const [engine, setEngine] = useState<MLCEngine | null>(null);
const [progressVisible, setProgressVisible] = useState(false);
const [progress, setProgress] = useState<Progress[]>([]);
const [loading, setLoading] = useState<boolean>(false);
const [sending, setSending] = useState<boolean>(false);

messages: Stores the chat history.
input: The current input from the user.
model: The selected model to load.
loadedModel: The currently loaded model.
engine: The MLCEngine instance that runs the LLM.
progress: Tracks the progress of model loading.
loading & sending: Flags to indicate if the app is currently loading a model or processing a message.

Refs for Scrolling

const progressContainerRef = useRef<HTMLDivElement>(null);
const chatBoxRef = useRef<HTMLDivElement>(null);

These refs are used to automatically scroll the chat and progress containers to the bottom as new content is added.

useEffect Hooks

useEffect(() => {
  scrollToBottom(chatBoxRef);
}, [messages]); // `messages` is your array containing the chat messages

useEffect(() => {
  scrollToBottom(progressContainerRef);
}, [progress]); // `progressLogs` is your array of progress logs

Scrolling Function

const scrollToBottom = (ref: RefObject<HTMLDivElement>) => {
  if (ref.current) {
    ref.current.scrollTop = ref.current.scrollHeight;
  }
};

This function manipulates the scroll position of the chat and progress containers.

Loading the Model

Progress Callback

const initProgressCallback = (initProgress: Progress) => {
  setProgress((prevProgress) => [...prevProgress, initProgress]);
};

This function updates the progress state as the model is being loaded.

Model Loading Function

const loadModel = async () => {
  if (loading === true) {
    alert("Already loading");
    return;
  }

  setLoading(true);

  try {
    const newEngine = await CreateMLCEngine(model, {
      initProgressCallback,
    });
    setEngine(newEngine);
    setLoadedModel(model);
  } catch (error) {
    alert(error.message);
  }

  setLoading(false);
};

CreateMLCEngine: Initializes the MLCEngine with the selected model.
initProgressCallback: Provides progress updates during model initialization.
Error Handling: Alerts the user if there's an error during loading.

Load Model Event Handler

const handleLoadModel = (e: React.FormEvent) => {
  e.preventDefault();
  setProgress([]);
  loadModel();
};

This handler resets the progress state and triggers the model loading process.

Handling User Input and Messages

Sending Messages

const handleSendMessage = async (e: React.FormEvent) => {
  e.preventDefault();

  if (!engine) {
    alert("Please load an engine before sending a message");
    return;
  }

  if (sending === true) {
    alert("A message is already being processed");
    return;
  }

  setSending(true);

  let newMessage: ChatMessage | null = null;

  // Create a new message or reuse the last one if input is empty
  if (input.trim()) {
    newMessage = { role: "user", message: input.trim() };
    setMessages((prevMessages) => [...prevMessages, newMessage]);
    setInput("");
  } else {
    // Reuse the last user message if input is empty
    for (let i = messages.length - 1; i >= 0; i--) {
      if (messages[i].role === "user") {
        newMessage = messages[i];
        setMessages((prevMessages) =>
          prevMessages.slice(0, -(messages.length - 1 - i))
        );
        break;
      }
    }
  }

  if (!newMessage) {
    setSending(false);
    return;
  }

  try {
    const botMessageChunks: AsyncIterable<ChatCompletionChunk> =
      await getBotResponse(
        engine,
        mapChatMessageToChatCompletionParam([...messages, newMessage])
      );

    // Add a placeholder for the assistant's response
    setMessages((prevMessages) => [
      ...prevMessages,
      { role: "assistant", message: "" },
    ]);

    // Process the streamed response
    for await (const chunk of botMessageChunks) {
      const chunkContent: string = chunk.choices[0]?.delta.content || "";

      if (chunkContent) {
        setMessages((prevMessages) => {
          const newMessages = [...prevMessages];
          const lastMessageIndex = newMessages.length - 1;
          newMessages[lastMessageIndex] = {
            ...newMessages[lastMessageIndex],
            message: newMessages[lastMessageIndex].message + chunkContent,
          };

          return newMessages;
        });
      }
    }
  } catch (error) {
    alert(error.message);
  }
  setSending(false);
};

Input Handling: Captures user input and adds it to the messages.
Engine Check: Ensures that a model is loaded before sending a message.
Streaming Responses: Processes the assistant's response in chunks, updating the UI in real-time.
Error Handling: Alerts the user if an error occurs during message processing.

Getting the Bot Response

const getBotResponse = async (
  engine: MLCEngine,
  messages: ChatCompletionMessageParam[]
): Promise<AsyncIterable<ChatCompletionChunk>> => {
  return await engine.chat.completions.create({
    messages,
    stream: true,
  });
};

This function interacts with the mlc-llm API to get the assistant's response as an asynchronous iterable, allowing us to process the response incrementally.

Rendering the Component

Chat Messages

<div className="chat__box" ref={chatBoxRef}>
  {messages.map((message, index) => (
    <div key={index} className={`chat__message chat__message--${message.role}`}>
      <ReactMarkdown>{message.message}</ReactMarkdown>
    </div>
  ))}
</div>

Message List: Iterates over the messages state to display each message.
Styling: Applies different styles based on the message role (assistant or user).
Markdown Support: Renders messages using ReactMarkdown to support Markdown syntax.

User Input Form

<form onSubmit={handleSendMessage} className="chat__form">
  <input
    type="text"
    value={input}
    onChange={(e)=> setInput(e.target.value)}
    placeholder="Type your message here..."
    className="chat__input"
  />
  <button
    type="submit"
    className="chat__button send"
    disabled={sending || loading}
  >
    Send
  </button>
</form>

Input Field: Captures user messages.
Send Button: Disabled while loading or sending to prevent duplicate actions.

Model Picker

<div className="model-picker">
  <select
    select
    onChange={(e)=> setModel(e.target.value)}
    id="model-select"
    className="chat__model_select"
  >
    {/* Options for different models */}
  </select>
</div>

Allows the user to select from a list of available models.

Loading and Progress Indicators

<button
  onClick={handleLoadModel}
  className="chat__button load"
  disabled={sending || loading || model === loadedModel}
>
  {loading ? "Loading..." : "Load Engine"}
</button>;
<progress
    className="loading-bar"
    hidden={!loading}
    value={progress?.[progress.length - 1]?.progress}
    max="1"
></progress>

Load Engine Button: Initiates the model loading process.
Progress Bar: Visual indicator of loading progress.

Progress Logs

<button
  type="button"
  onClick={()=>
    setProgressVisible((prevProgressVisible)=> !prevProgressVisible)
  }
  className="progress__button"
  disabled={!progress.length}
>
  {progressVisible ? "Hide Progress" : "Show Progress"}
</button>;

<div
  className="progress__container"
  hidden={!progressVisible}
  ref={progressContainerRef}
>
  {progress.map((step, index) => (
    <p key={index} className="progress__log">
      - {step.text}
    </p>
  ))}
</div>;

Show/Hide Progress Button: Toggles the visibility of the progress logs.
Progress Container: Displays detailed progress information during model loading and message processing.

Integrating MLC-LLM and WebGPU

The mlc-llm package is designed to run quantized LLMs efficiently in the browser by leveraging WebGPU. Here's how the integration works in our component:

Model Loading: When CreateMLCEngine is called, it initializes the engine with the selected model, which is loaded into the browser's memory.
WebGPU Utilization: The heavy computations required by the LLM are offloaded to the GPU via WebGPU, providing significant performance improvements over traditional CPU-bound operations.
Asynchronous Processing: Both model loading and message processing are handled asynchronously to keep the UI responsive.
Streaming Responses: By processing the assistant's response in chunks, we can provide real-time feedback to the user as the model generates text.

Working with Small LLMs

Using smaller LLMs in the browser has several advantages:

Performance: Smaller models require less computational power, making them suitable for client-side execution.
Privacy: User data doesn't need to be sent to a server, enhancing privacy.
Interactivity: Immediate feedback and reduced latency improve the user experience.

In our application, we offer a selection of small LLMs that the user can choose from, balancing performance and capability according to their needs.

Conclusion

By combining React, mlc-llm, and WebGPU, we've created a powerful and interactive chat application that runs entirely in the browser. This approach opens up new possibilities for web-based AI applications, providing users with responsive and private interactions with LLMs.

Key Takeaways

WebGPU enables efficient on-device computations in the browser.
mlc-llm provides a convenient interface for running LLMs client-side.
Asynchronous programming and streaming responses are crucial for a responsive UI.
State management and effects in React allow for smooth user interactions.

Next Steps

Expand Model Selection: Explore more models and experiment with different sizes and capabilities.
Enhance UI/UX: Improve the chat interface with better styling and additional features like conversation history.
Performance Optimization: Optimize loading times and response generation for even smoother interactions.

References

Feel free to clone the code and experiment with building your own browser-based AI applications!

Integrating Small LLMs Into Webpages with WebGPU Using MLC-LLM