Building In-Browser AI Applications

npm install @mlc-ai/web-llm

import * as webllm from "@mlc-ai/web-llm";

const initProgressCallback = (progress) => {
  // progress.text contains a human-readable status message
  // progress.progress is 0-1 for download progress
  document.getElementById("status").textContent = progress.text;

  if (progress.progress !== undefined) {
    document.getElementById("progress-bar").style.width =
      `${Math.round(progress.progress * 100)}%`;
  }
};

const engine = await webllm.CreateMLCEngine(
  "gemma-2-2b-it-q4f16_1-MLC", // Model identifier
  {
    initProgressCallback,
    logLevel: "INFO",
  }
);

const response = await engine.chat.completions.create({
  messages: [
    {
      role: "system",
      content: "You are a helpful assistant that summarises documents concisely."
    },
    {
      role: "user",
      content: `Summarise the following document:\n\n${documentText}`
    }
  ],
  temperature: 0.3,
  max_tokens: 500,
});

const summary = response.choices[0].message.content;

const chunks = await engine.chat.completions.create({
  messages: [
    { role: "user", content: "Explain the key risks in this contract." }
  ],
  temperature: 0.3,
  max_tokens: 800,
  stream: true,
});

let fullResponse = "";
for await (const chunk of chunks) {
  const delta = chunk.choices[0]?.delta?.content || "";
  fullResponse += delta;
  document.getElementById("output").textContent = fullResponse;
}

async function initAI() {
  // Check WebGPU support
  if (!navigator.gpu) {
    showFallbackMessage("Your browser does not support WebGPU. " +
      "Please use Chrome 113+ or Safari 17+.");
    return null;
  }

  // Check GPU adapter availability
  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) {
    showFallbackMessage("No GPU adapter found. " +
      "Your device may not have a compatible GPU.");
    return null;
  }

  // Check available memory (heuristic)
  const adapterInfo = await adapter.requestAdapterInfo();
  console.log("GPU:", adapterInfo.description);

  try {
    const engine = await webllm.CreateMLCEngine(
      "gemma-2-2b-it-q4f16_1-MLC",
      { initProgressCallback }
    );
    return engine;
  } catch (error) {
    if (error.message.includes("out of memory")) {
      showFallbackMessage("Insufficient GPU memory for this model. " +
        "Try closing other tabs or applications.");
    } else {
      showFallbackMessage(`Failed to load AI model: ${error.message}`);
    }
    return null;
  }
}

import { pipeline, env } from '@huggingface/transformers';

// Configure caching
env.cacheDir = './.cache';
env.allowLocalModels = false;

// Create a text generation pipeline
const generator = await pipeline(
  'text-generation',
  'onnx-community/gemma-2-2b-it-ONNX', // ONNX-format model
  {
    device: 'webgpu',        // Use WebGPU; falls back to 'wasm' automatically
    dtype: 'q4',             // INT4 quantisation
  }
);

// Generate with streaming via a callback
const output = await generator(
  [
    { role: "user", content: "Summarise this document briefly." }
  ],
  {
    max_new_tokens: 256,
    temperature: 0.3,
    do_sample: true,
    callback_function: (token) => {
      // Called for each generated token
      process.stdout.write(token);
    }
  }
);

// Embedding pipeline (for RAG or search)
const embedder = await pipeline(
  'feature-extraction',
  'Xenova/nomic-embed-text-v1.5',
  { device: 'webgpu', dtype: 'fp16' }
);

// Embed a document
const embedding = await embedder("Your document text here", {
  pooling: 'mean',
  normalize: true
});

// embedding.data is a Float32Array of 768 dimensions

// Strategy 1: Check available memory before loading
async function estimateAvailableGPUMemory() {
  const adapter = await navigator.gpu?.requestAdapter();
  if (!adapter) return 0;

  // Request device with maximum limits
  const device = await adapter.requestDevice({
    requiredLimits: {
      maxBufferSize: adapter.limits.maxBufferSize,
      maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
    }
  });

  // maxBufferSize gives an upper bound on single allocation
  // Total available is typically 2-4x this
  return adapter.limits.maxBufferSize;
}

// Strategy 2: Graceful degradation
async function loadModelWithFallback() {
  const availableMemory = await estimateAvailableGPUMemory();

  if (availableMemory > 4 * 1024 * 1024 * 1024) {
    // 4GB+ available: try E4B
    return await tryLoadModel("gemma-4-e4b-q4");
  } else if (availableMemory > 1.5 * 1024 * 1024 * 1024) {
    // 1.5GB+ available: use E2B
    return await tryLoadModel("gemma-4-e2b-q4");
  } else {
    // Not enough GPU memory: fall back to WASM
    return await tryLoadModel("gemma-4-e2b-q4", { device: "wasm" });
  }
}

// Strategy 3: Release memory when not in use
function releaseModel(engine) {
  engine.unload(); // Releases GPU buffers
  // Model will need to be reloaded from IndexedDB cache for next use
}

// Production chat implementation with WebLLM
class BrowserChat {
  constructor() {
    this.engine = null;
    this.conversationHistory = [];
    this.isGenerating = false;
  }

  async initialise(statusCallback) {
    statusCallback({ phase: "checking", message: "Checking GPU support..." });

    if (!navigator.gpu) {
      statusCallback({ phase: "error", message: "WebGPU not supported" });
      return false;
    }

    statusCallback({ phase: "loading", message: "Loading AI model..." });

    try {
      this.engine = await webllm.CreateMLCEngine(
        "gemma-2-2b-it-q4f16_1-MLC",
        {
          initProgressCallback: (progress) => {
            statusCallback({
              phase: "loading",
              message: progress.text,
              progress: progress.progress
            });
          }
        }
      );

      // Warm-up inference
      statusCallback({ phase: "warming", message: "Warming up..." });
      await this.engine.chat.completions.create({
        messages: [{ role: "user", content: "Hi" }],
        max_tokens: 1,
      });

      statusCallback({ phase: "ready", message: "AI ready" });
      return true;

    } catch (error) {
      statusCallback({
        phase: "error",
        message: `Failed to load: ${error.message}`
      });
      return false;
    }
  }

  async sendMessage(userMessage, onToken) {
    if (this.isGenerating) return;
    this.isGenerating = true;

    this.conversationHistory.push({
      role: "user",
      content: userMessage
    });

    try {
      const stream = await this.engine.chat.completions.create({
        messages: [
          {
            role: "system",
            content: "You are a helpful enterprise assistant. Be concise and accurate."
          },
          ...this.conversationHistory
        ],
        temperature: 0.3,
        max_tokens: 1024,
        stream: true,
      });

      let assistantMessage = "";
      for await (const chunk of stream) {
        const delta = chunk.choices[0]?.delta?.content || "";
        assistantMessage += delta;
        onToken(delta, assistantMessage);
      }

      this.conversationHistory.push({
        role: "assistant",
        content: assistantMessage
      });

    } catch (error) {
      onToken(null, null, error);
    } finally {
      this.isGenerating = false;
    }
  }

  resetConversation() {
    this.conversationHistory = [];
    this.engine.resetChat();
  }
}