Spaces:

webml-community
/

Baguettotron-WebGPU

Running

App Files Files Community

Xenova HF Staff commited on 15 days ago

Commit

c369b38

verified ·

1 Parent(s): c1521a5

Upload 10 files

Browse files

Files changed (10) hide show

eslint.config.js +23 -0
index.html +16 -0
package.json +35 -0
src/App.tsx +763 -0
src/index.css +54 -0
src/main.tsx +10 -0
tsconfig.app.json +28 -0
tsconfig.json +7 -0
tsconfig.node.json +26 -0
vite.config.ts +8 -0

eslint.config.js ADDED Viewed

	@@ -0,0 +1,23 @@

+import js from "@eslint/js";
+import globals from "globals";
+import reactHooks from "eslint-plugin-react-hooks";
+import reactRefresh from "eslint-plugin-react-refresh";
+import tseslint from "typescript-eslint";
+import { defineConfig, globalIgnores } from "eslint/config";
+export default defineConfig([
+  globalIgnores(["dist"]),
+  {
+    files: ["**/*.{ts,tsx}"],
+    extends: [
+      js.configs.recommended,
+      tseslint.configs.recommended,
+      reactHooks.configs.flat.recommended,
+      reactRefresh.configs.vite,
+    ],
+    languageOptions: {
+      ecmaVersion: 2020,
+      globals: globals.browser,
+    },
+  },
+]);

index.html ADDED Viewed

	@@ -0,0 +1,16 @@

+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <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>"
+    />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>Baguettotron WebGPU</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/main.tsx"></script>
+  </body>
+</html>

package.json ADDED Viewed

	@@ -0,0 +1,35 @@

+{
+  "name": "baguettotron-webgpu",
+  "private": true,
+  "version": "0.0.0",
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "tsc -b && vite build",
+    "lint": "eslint .",
+    "preview": "vite preview"
+  },
+  "dependencies": {
+    "@huggingface/transformers": "^3.7.6",
+    "@tailwindcss/vite": "^4.1.17",
+    "lucide-react": "^0.553.0",
+    "react": "^19.2.0",
+    "react-dom": "^19.2.0",
+    "streamdown": "^1.4.0",
+    "tailwindcss": "^4.1.17"
+  },
+  "devDependencies": {
+    "@eslint/js": "^9.39.1",
+    "@types/node": "^24.10.0",
+    "@types/react": "^19.2.2",
+    "@types/react-dom": "^19.2.2",
+    "@vitejs/plugin-react": "^5.1.0",
+    "eslint": "^9.39.1",
+    "eslint-plugin-react-hooks": "^7.0.1",
+    "eslint-plugin-react-refresh": "^0.4.24",
+    "globals": "^16.5.0",
+    "typescript": "~5.9.3",
+    "typescript-eslint": "^8.46.3",
+    "vite": "^7.2.2"
+  }
+}

src/App.tsx ADDED Viewed

	@@ -0,0 +1,763 @@

+import { useState, useRef, useCallback, useLayoutEffect } from "react";
+import { Send, Paperclip, Brain, ChevronDown, X, Plus } from "lucide-react";
+import {
+  AutoModelForCausalLM,
+  AutoTokenizer,
+  InterruptableStoppingCriteria,
+  TextStreamer,
+} from "@huggingface/transformers";
+import { Streamdown } from "streamdown";
+import type {
+  PreTrainedTokenizer,
+  LlamaForCausalLM,
+} from "@huggingface/transformers";
+import type React from "react";
+const MODEL_ID = "onnx-community/Baguettotron-ONNX";
+const DTYPES = {
+  fp32: "FP32 (~1.28 GB)",
+  fp16: "FP16 (~642 MB)",
+  q4: "Q4 (~329 MB)",
+  q4f16: "Q4F16 (~235 MB)",
+} as const;
+type Dtype = keyof typeof DTYPES;
+const SOURCE_SEPARATOR_REGEX = /\n{2,}/g;
+type Role = "user" | "assistant";
+/**
+ * Format the sources into tagged segments for the model input.
+ */
+const buildSourcesPayload = (rawContext: string) => {
+  const trimmed = rawContext.trim();
+  if (!trimmed) {
+    return { payload: "", count: 0, segments: [] };
+  }
+  const segments = trimmed
+    .split(SOURCE_SEPARATOR_REGEX)
+    .map((segment) => segment.trim())
+    .filter(Boolean);
+  if (segments.length === 0) {
+    return { payload: "", count: 0, segments: [] };
+  }
+  const payload =
+    "\n\n" +
+    segments
+      .map(
+        (segment, index) =>
+          `<source_${index + 1}>${segment}</source_${index + 1}>`,
+      )
+      .join("\n");
+  return { payload, count: segments.length, segments };
+};
+/**
+ * Converts <ref>...</ref> tags in the content to superscript references.
+ */
+const convertRefsToSuperscript = (content: string) => {
+  const refRegex = /<ref name="([^"]+)">([\s\S]*?)<\/ref>/g;
+  const refLabelMap = new Map<string, number>();
+  let refCounter = 1;
+  // First, process all complete <ref>...</ref> tags
+  let result = content.replace(refRegex, (_, sourceName = "", refBody) => {
+    const label =
+      refLabelMap.get(sourceName) ??
+      (() => {
+        const assigned = refCounter++;
+        refLabelMap.set(sourceName, assigned);
+        return assigned;
+      })();
+    const escapedRefBody = refBody.replace(/"/g, "&quot;");
+    return `<sup className="cursor-pointer" title="${escapedRefBody}">[${label}]</sup>`;
+  });
+  // Remove any trailing incomplete <ref> tag
+  const incompleteRefRegex = /<ref[^>]*>[\s\S]*$/;
+  result = result.replace(incompleteRefRegex, "");
+  return result;
+};
+/**
+ * Sanitizes user input by replacing angle brackets.
+ */
+const sanitizeInput = (text: string) => {
+  return text.replace(/</g, "&lt;").replace(/>/g, "&gt;");
+};
+/**
+ * Represents a single chat message in the history.
+ */
+interface Message {
+  id: number;
+  role: Role;
+  content: string;
+  thinkTrace?: string;
+  rawStream?: string;
+  isLoading?: boolean;
+  timestamp?: number;
+  thinkEndTime?: number;
+}
+/**
+ * A simple, self-contained collapsible component.
+ */
+const Collapsible: React.FC<{
+  title: React.ReactNode;
+  children: React.ReactNode;
+}> = ({ title, children }) => {
+  const [isOpen, setIsOpen] = useState(false);
+  const contentRef = useRef<HTMLDivElement>(null);
+  return (
+    <div className="collapsible mt-2">
+      <button
+        onClick={() => setIsOpen(!isOpen)}
+        className="flex items-center space-x-1 text-xs font-medium text-amber-700 hover:text-amber-900 transition-colors"
+      >
+        {title}
+        <ChevronDown
+          size={14}
+          className={`transform transition-transform ${isOpen ? "rotate-180" : "rotate-0"}`}
+        />
+      </button>
+      <div
+        ref={contentRef}
+        style={{
+          maxHeight: isOpen ? `${contentRef.current?.scrollHeight}px` : "0px",
+        }}
+        className="overflow-hidden transition-all duration-300 ease-in-out"
+      >
+        <div className="mt-2 p-2 bg-amber-50 border border-dashed border-amber-200 rounded-md text-xs text-stone-600 prose-sm">
+          {children}
+        </div>
+      </div>
+    </div>
+  );
+};
+/**
+ * A single chat message bubble.
+ */
+const MessageBubble: React.FC<{ message: Message; minHeight?: number }> = ({
+  message,
+  minHeight,
+}) => {
+  const { role, content, thinkTrace, isLoading, timestamp, thinkEndTime } =
+    message;
+  const isUser = role === "user";
+  let thinkingText = "";
+  let opacityClass = "";
+  const hasDuration =
+    typeof thinkEndTime === "number" && typeof timestamp === "number";
+  const durationSeconds = hasDuration
+    ? Math.max(Math.round((thinkEndTime - timestamp) / 1000), 0)
+    : null;
+  if (isLoading && !thinkEndTime) {
+    thinkingText = "Thinking...";
+    opacityClass = "opacity-70 hover:opacity-100";
+  } else if (thinkTrace) {
+    thinkingText =
+      durationSeconds !== null
+        ? `Thought for ${durationSeconds} seconds`
+        : "Thought interrupted";
+  } else {
+    thinkingText = "Show Thoughts";
+  }
+  const markdownContent = convertRefsToSuperscript(content);
+  return (
+    <div
+      data-message-id={message.id}
+      data-role={role}
+      className={`message flex items-start animate-in fade-in slide-in-from-bottom-2 duration-300 py-2 ${isUser ? "justify-end" : "justify-start"}`}
+      style={{
+        minHeight,
+      }}
+    >
+      <div
+        className={`max-w-xl lg:max-w-2xl px-4 py-3 rounded-2xl ${
+          isUser
+            ? "bg-amber-500 text-white rounded-br-none"
+            : "bg-white text-stone-800 rounded-bl-none shadow-sm border border-stone-200"
+        }`}
+      >
+        {(thinkTrace || isLoading) && (
+          <Collapsible
+            title={
+              <div className="flex items-center space-x-1.5 text-sm">
+                <Brain size={16} />
+                <span
+                  className={`${isLoading ? "animate-glisten" : ""} ${opacityClass}`}
+                >
+                  {thinkingText}
+                </span>
+              </div>
+            }
+          >
+            <Streamdown
+              parseIncompleteMarkdown={false}
+              className="text-xs text-stone-500"
+              isAnimating={Boolean(isLoading && thinkEndTime)}
+            >
+              {thinkTrace || (isLoading ? "..." : "")}
+            </Streamdown>
+          </Collapsible>
+        )}
+        <div className={`${thinkTrace || isLoading ? "mt-2" : ""}`}>
+          <Streamdown
+            parseIncompleteMarkdown={false}
+            className="text-sm leading-relaxed text-stone-800"
+            isAnimating={Boolean(isLoading && !thinkEndTime)}
+          >
+            {markdownContent || (isLoading ? "" : "")}
+          </Streamdown>
+        </div>
+      </div>
+    </div>
+  );
+};
+/**
+ * Manages the model and tokenizer loading state and refs.
+ */
+const useLLM = () => {
+  const [modelStatus, setModelStatus] = useState<
+    "idle" | "loading" | "ready" | "error"
+  >("idle");
+  const [loadProgress, setLoadProgress] = useState(0);
+  const modelRef = useRef<LlamaForCausalLM | null>(null);
+  const tokenizerRef = useRef<PreTrainedTokenizer | null>(null);
+  const loadModel = useCallback(
+    async (dtype: Dtype) => {
+      if (modelRef.current && tokenizerRef.current) {
+        setModelStatus("ready");
+        setLoadProgress(100);
+        return;
+      }
+      if (modelStatus === "loading") return;
+      setModelStatus("loading");
+      setLoadProgress(0);
+      const progress_callback = (progress: any) => {
+        if (
+          progress.status === "progress" &&
+          typeof progress.total === "number" &&
+          typeof progress.loaded === "number" &&
+          typeof progress.file === "string" &&
+          progress.file.endsWith(".onnx_data")
+        ) {
+          const percentage = Math.round(
+            (progress.loaded / progress.total) * 100,
+          );
+          setLoadProgress(percentage);
+        }
+      };
+      try {
+        const tokenizer = await AutoTokenizer.from_pretrained(MODEL_ID, {
+          progress_callback,
+        });
+        const model = await AutoModelForCausalLM.from_pretrained(MODEL_ID, {
+          dtype,
+          device: "webgpu",
+          progress_callback,
+        });
+        tokenizerRef.current = tokenizer;
+        modelRef.current = model;
+        setLoadProgress(100);
+        setModelStatus("ready");
+      } catch (error) {
+        console.error("Failed to load model", error);
+        setModelStatus("error");
+      }
+    },
+    [modelStatus],
+  );
+  return {
+    modelStatus,
+    loadProgress,
+    modelRef,
+    tokenizerRef,
+    loadModel,
+  };
+};
+const App: React.FC = () => {
+  const [messages, setMessages] = useState<Message[]>([]);
+  const [currentInput, setCurrentInput] = useState("");
+  const [context, setContext] = useState("");
+  const [showContext, setShowContext] = useState(false);
+  const [isLoading, setIsLoading] = useState(false);
+  const [lastMessageMinHeight, setLastMessageMinHeight] = useState<
+    number | undefined
+  >(undefined);
+  const [selectedDtype, setSelectedDtype] = useState<Dtype>("fp16");
+  const [dtypeMenuOpen, setDtypeMenuOpen] = useState(false);
+  const stoppingCriteriaRef = useRef<InterruptableStoppingCriteria | null>(
+    null,
+  );
+  const mainRef = useRef<HTMLDivElement>(null);
+  const { modelStatus, loadProgress, modelRef, tokenizerRef, loadModel } =
+    useLLM();
+  useLayoutEffect(() => {
+    if (!mainRef.current) return;
+    const el = mainRef.current;
+    // If the last message is from the assistant, calculate a min-height to prevent layout shifts.
+    if (messages.at(-1)?.role === "assistant") {
+      const userMessageElement = el.querySelector<HTMLDivElement>(
+        `[data-message-id="${messages.at(-2)?.id}"]`,
+      );
+      if (userMessageElement) {
+        const userMessageHeight =
+          userMessageElement.getBoundingClientRect().height;
+        const screenHeight = window.innerHeight;
+        const newMinHeight = Math.max(
+          screenHeight - userMessageHeight - 270,
+          0,
+        );
+        setLastMessageMinHeight(newMinHeight);
+      }
+    } else {
+      setLastMessageMinHeight(undefined);
+    }
+  }, [messages.length]);
+  useLayoutEffect(() => {
+    if (mainRef.current) {
+      const el = mainRef.current;
+      setTimeout(() => {
+        el.scrollTo({
+          top: el.scrollHeight,
+          behavior: "smooth",
+        });
+      }, 0);
+    }
+  }, [messages.length, lastMessageMinHeight]);
+  const handleStreamUpdate = useCallback((newToken: string) => {
+    setMessages((prev) => {
+      if (prev.length === 0 || prev.at(-1)!.role === "user") {
+        return prev;
+      }
+      const lastMessage = { ...prev.at(-1)! };
+      lastMessage.rawStream = (lastMessage.rawStream || "") + newToken;
+      const raw = lastMessage.rawStream;
+      const thinkEndTag = "</think>";
+      const thinkEndIndex = raw.indexOf(thinkEndTag);
+      let content;
+      let thinkTrace = "";
+      if (thinkEndIndex !== -1) {
+        // Think block is complete.
+        thinkTrace = raw.substring(0, thinkEndIndex);
+        const contentAfter = raw.substring(thinkEndIndex + thinkEndTag.length);
+        content = contentAfter.replace("<|im_end|><|end_of_text|>", "");
+        if (!lastMessage.thinkEndTime) {
+          lastMessage.thinkEndTime = Date.now();
+        }
+      } else {
+        // Think block has started but not finished.
+        thinkTrace = raw;
+        content = "";
+      }
+      lastMessage.content = content.trim();
+      lastMessage.thinkTrace = thinkTrace.trim();
+      return [...prev.slice(0, -1), lastMessage];
+    });
+  }, []);
+  const handleStopGeneration = useCallback(() => {
+    stoppingCriteriaRef.current?.interrupt();
+  }, []);
+  const streamAssistantResponse = useCallback(
+    async (
+      historyForModel: { role: Role; content: string }[],
+      assistantMessageId: number,
+    ) => {
+      const tokenizer = tokenizerRef.current;
+      const model = modelRef.current;
+      if (!tokenizer || !model) return;
+      const inputs = tokenizer.apply_chat_template(historyForModel, {
+        add_generation_prompt: true,
+        return_dict: true,
+      }) as any;
+      const streamer = new TextStreamer(tokenizer, {
+        skip_prompt: true,
+        skip_special_tokens: false,
+        callback_function: (token: string) => handleStreamUpdate(token),
+      });
+      const stoppingCriteria = new InterruptableStoppingCriteria();
+      stoppingCriteriaRef.current = stoppingCriteria;
+      try {
+        await model.generate({
+          ...inputs,
+          max_new_tokens: 2048,
+          streamer,
+          stopping_criteria: stoppingCriteria,
+          repetition_penalty: 1.2,
+        });
+      } catch (error) {
+        console.error(error);
+      } finally {
+        stoppingCriteriaRef.current = null;
+        setIsLoading(false);
+        setMessages((prev) =>
+          prev.map((msg) => {
+            if (msg.id === assistantMessageId) {
+              const { rawStream, isLoading: _, ...rest } = msg;
+              return rest;
+            }
+            return msg;
+          }),
+        );
+      }
+    },
+    [handleStreamUpdate, modelRef, tokenizerRef],
+  );
+  const handleSubmit = async (
+    e?: React.FormEvent,
+    prompt?: string,
+    sources?: string,
+  ) => {
+    if (e) e.preventDefault();
+    if (isLoading || modelStatus !== "ready") return;
+    const input = prompt || currentInput;
+    if (!input.trim()) return;
+    const trimmedContext = (sources || context).trim();
+    const {
+      payload: sourcesPayload,
+      count: sourceCount,
+      segments: sourceSegments,
+    } = buildSourcesPayload(trimmedContext);
+    const fullPrompt = `${input}${sourcesPayload}`;
+    const sanitizedInput = sanitizeInput(input);
+    let userMessageContent = sanitizedInput;
+    if (sourceCount > 0) {
+      const sourcesList = sourceSegments
+        .map(
+          (seg, i) =>
+            `${i + 1}. ${seg.substring(0, 75)}${seg.length > 75 ? "..." : ""}`,
+        )
+        .join("\n");
+      userMessageContent += `\n\n[Source${sourceCount > 1 ? "s" : ""}]:\n${sourcesList}`;
+    }
+    const userMessage: Message = {
+      id: messages.length,
+      role: "user",
+      content: userMessageContent,
+    };
+    const assistantPlaceholder: Message = {
+      id: messages.length + 1,
+      role: "assistant",
+      content: "",
+      thinkTrace: "",
+      rawStream: "",
+      isLoading: true,
+      timestamp: Date.now(),
+    };
+    setMessages((prev) => [...prev, userMessage, assistantPlaceholder]);
+    setCurrentInput("");
+    setContext("");
+    setShowContext(false);
+    setIsLoading(true);
+    setLastMessageMinHeight(undefined);
+    const historyForModel = [
+      ...messages.map(({ role, content }) => ({ role, content })),
+      { role: "user" as Role, content: fullPrompt },
+    ];
+    await streamAssistantResponse(historyForModel, assistantPlaceholder.id);
+  };
+  const handleNewChat = () => {
+    handleStopGeneration();
+    setMessages([]);
+    setCurrentInput("");
+    setContext("");
+    setShowContext(false);
+    setIsLoading(false);
+    setLastMessageMinHeight(undefined);
+  };
+  return (
+    <div className="flex flex-col h-screen bg-amber-50 font-sans text-stone-800">
+      {modelStatus === "ready" && (
+        <header className="flex-shrink-0 sticky top-0 z-10 flex items-center justify-between p-4 bg-white/90 backdrop-blur-md shadow-sm border-b border-amber-200 h-[100px]">
+          <button
+            onClick={handleNewChat}
+            className="p-2 rounded-full text-stone-500 hover:text-amber-600 hover:bg-amber-50 transition-colors"
+            title="New Chat"
+          >
+            <Plus size={20} />
+          </button>
+          <div className="flex-1 text-center">
+            <h1 className="text-2xl md:text-3xl font-serif font-bold text-amber-800">
+              🥖 Baguettotron WebGPU
+            </h1>
+            <p className="text-sm text-stone-600">
+              A small but powerful reasoning model
+            </p>
+          </div>
+        </header>
+      )}
+      <main ref={mainRef} className="flex-grow overflow-y-auto">
+        <div className="mx-auto w-full max-w-6xl p-4 md:p-6 space-y-2 h-full">
+          {modelStatus !== "ready" ? (
+            <div className="flex h-full flex-col items-center justify-center gap-6 text-center text-stone-600">
+              <span className="text-8xl animate-wobble">🥖</span>
+              <div>
+                <h1 className="text-5xl font-bold text-amber-800">
+                  Baguettotron WebGPU
+                </h1>
+                <p className="mt-4 max-w-xl text-md">
+                  You are about to load Baguettotron, a 300M parameter reasoning
+                  model optimized for in-browser inference. Everything runs
+                  entirely in your browser with 🤗 Transformers.js and ONNX
+                  Runtime Web, meaning no data is sent to a server. Once loaded,
+                  it can even be used offline.
+                </p>
+              </div>
+              <div className="relative inline-flex rounded-full shadow-sm">
+                <button
+                  onClick={() => loadModel(selectedDtype)}
+                  disabled={modelStatus === "loading"}
+                  className="rounded-l-full bg-amber-600 pl-6 pr-5 py-3 text-white font-medium transition hover:bg-amber-700 disabled:opacity-50 disabled:cursor-not-allowed"
+                >
+                  {modelStatus === "loading"
+                    ? `Loading ${loadProgress}%`
+                    : `Load model (${selectedDtype})`}
+                </button>
+                <button
+                  onClick={() => setDtypeMenuOpen(!dtypeMenuOpen)}
+                  disabled={modelStatus === "loading"}
+                  className="rounded-r-full bg-amber-600 px-3 py-3 text-white transition hover:bg-amber-700 disabled:opacity-50 border-l border-amber-500"
+                >
+                  <ChevronDown
+                    size={20}
+                    className={`transform transition-transform ${dtypeMenuOpen ? "rotate-180" : ""}`}
+                  />
+                </button>
+                {dtypeMenuOpen && (
+                  <div className="absolute top-full mt-2 w-full bg-white rounded-md shadow-lg z-10 border border-stone-200">
+                    {Object.entries(DTYPES).map(([dtype, label]) => (
+                      <button
+                        key={dtype}
+                        onClick={() => {
+                          setSelectedDtype(dtype as Dtype);
+                          setDtypeMenuOpen(false);
+                        }}
+                        className="w-full text-left px-4 py-2 text-sm text-stone-700 hover:bg-amber-50"
+                      >
+                        {label}
+                      </button>
+                    ))}
+                  </div>
+                )}
+              </div>
+              {modelStatus === "error" && (
+                <p className="text-sm text-red-600">
+                  Model load failed. Check console for details and retry.
+                </p>
+              )}
+            </div>
+          ) : (
+            <>
+              {messages.length === 0 && (
+                <div className="flex flex-col items-center justify-center h-full text-center text-stone-500">
+                  <div className="p-8 rounded-2xl flex flex-col items-center">
+                    <h2 className="text-3xl font-semibold mt-4 text-stone-700">
+                      Welcome to Baguettotron
+                    </h2>
+                    <h3 className="max-w-xs mt-1 text-lg">
+                      Ask me a question, or try one of the examples below!
+                    </h3>
+                  </div>
+                  <div className="mt-2 flex flex-wrap justify-center gap-4">
+                    <button
+                      onClick={() =>
+                        handleSubmit(
+                          undefined,
+                          "What is the capital of France? Just provide the answer.",
+                        )
+                      }
+                      className="bg-amber-100 hover:bg-amber-200 text-amber-800 px-4 py-2 rounded-lg shadow-sm border border-amber-200 transition-colors"
+                    >
+                      Encyclopedic knowledge
+                    </button>
+                    {["fp32", "fp16"].includes(selectedDtype) && (
+                      <button
+                        onClick={() =>
+                          handleSubmit(
+                            undefined,
+                            "Write me a short poem about machine learning.",
+                          )
+                        }
+                        className="bg-amber-100 hover:bg-amber-200 text-amber-800 px-4 py-2 rounded-lg shadow-sm border border-amber-200 transition-colors"
+                      >
+                        Creative writing
+                      </button>
+                    )}
+                    <button
+                      onClick={() =>
+                        handleSubmit(
+                          undefined,
+                          "Which is wider: Australia or the Moon?",
+                          "Australia is approximately 4,000 km in width from east to west, according to Geoscience Australia.\n\nThe diameter of the Moon is about 3,476 km, according to Britannica.",
+                        )
+                      }
+                      className="bg-amber-100 hover:bg-amber-200 text-amber-800 px-4 py-2 rounded-lg shadow-sm border border-amber-200 transition-colors"
+                    >
+                      RAG with grounding
+                    </button>
+                  </div>
+                </div>
+              )}
+              {messages.map((msg, index) => {
+                const isLastAssistantMessage =
+                  index === messages.length - 1 && msg.role === "assistant";
+                const minHeight = isLastAssistantMessage
+                  ? lastMessageMinHeight
+                  : undefined;
+                return (
+                  <MessageBubble
+                    key={msg.id}
+                    message={msg}
+                    minHeight={minHeight}
+                  />
+                );
+              })}
+            </>
+          )}
+        </div>
+      </main>
+      {modelStatus === "ready" && (
+        <footer className="flex-shrink-0 sticky bottom-0 z-10 p-4 bg-white/90 backdrop-blur-md border-t border-amber-100">
+          <form onSubmit={handleSubmit} className="max-w-3xl mx-auto">
+            <div
+              style={{
+                maxHeight: showContext ? "120px" : "0px",
+                transition: "max-height 0.3s ease-in-out",
+                opacity: showContext ? 1 : 0,
+              }}
+              className="overflow-hidden relative"
+            >
+              <textarea
+                value={context}
+                onChange={(e) => setContext(e.target.value)}
+                disabled={isLoading}
+                placeholder="Add RAG context here. Separate multiple sources with two new lines."
+                className="w-full h-28 p-2 mb-2 rounded-lg border border-stone-300 focus:ring-amber-500 focus:border-amber-500 text-sm resize-none"
+              />
+              <button
+                type="button"
+                onClick={() => setShowContext(false)}
+                className="absolute top-2 right-2 p-1 text-stone-400 hover:text-stone-600 bg-white/50 rounded-full"
+              >
+                <X size={16} />
+              </button>
+            </div>
+            <div className="flex items-center space-x-2">
+              <button
+                type="button"
+                onClick={() => setShowContext(!showContext)}
+                title="Add Context for RAG"
+                className={`flex-shrink-0 p-2 rounded-full transition-colors ${
+                  showContext
+                    ? "bg-amber-100 text-amber-700"
+                    : "text-stone-500 hover:text-amber-600 hover:bg-amber-50"
+                }`}
+              >
+                <Paperclip size={20} />
+              </button>
+              <input
+                type="text"
+                value={currentInput}
+                onChange={(e) => setCurrentInput(e.target.value)}
+                placeholder="Send a message..."
+                className="flex-grow px-4 py-2 rounded-full border border-stone-300 focus:ring-2 focus:ring-amber-500 focus:border-transparent outline-none transition-shadow"
+                disabled={isLoading || modelStatus !== "ready"}
+              />
+              {isLoading ? (
+                <button
+                  type="button"
+                  onClick={handleStopGeneration}
+                  className="group flex h-10 w-10 flex-shrink-0 items-center justify-center rounded-full border border-stone-300 bg-white text-stone-600 hover:border-red-500"
+                >
+                  <span className="h-3.5 w-3.5 rounded-sm bg-stone-600 transition-colors group-hover:bg-red-500" />
+                </button>
+              ) : (
+                <button
+                  type="submit"
+                  disabled={
+                    isLoading || !currentInput.trim() || modelStatus !== "ready"
+                  }
+                  className="flex h-10 w-10 flex-shrink-0 items-center justify-center rounded-full bg-amber-600 text-white transition-all transform
+                             hover:bg-amber-700 hover:scale-105 active:scale-95
+                             disabled:bg-stone-300 disabled:scale-100 disabled:cursor-not-allowed"
+                >
+                  <Send size={20} />
+                </button>
+              )}
+            </div>
+            <p className="text-center text-xs text-stone-400 mt-2">
+              ⚡ Powered by{" "}
+              <a
+                href="https://github.com/huggingface/transformers.js"
+                target="_blank"
+                rel="noopener noreferrer"
+              >
+                Transformers.js
+              </a>{" "}
+              — Runs locally in your browser on WebGPU.
+            </p>
+          </form>
+        </footer>
+      )}
+    </div>
+  );
+};
+export default App;

src/index.css ADDED Viewed

	@@ -0,0 +1,54 @@

+@import "tailwindcss";
+@source "../node_modules/streamdown/dist/index.js";
+@keyframes glisten {
+  0% {
+    background-position: 200% 0;
+  }
+  100% {
+    background-position: -200% 0;
+  }
+}
+.animate-glisten {
+  /* Using amber-700 and amber-300 for the gradient */
+  background-image: linear-gradient(
+    90deg,
+    rgba(168, 85, 6, 1) 0%,
+    rgba(253, 186, 116, 1) 40%,
+    rgba(253, 186, 116, 1) 60%,
+    rgba(168, 85, 6, 1) 100%
+  );
+  color: transparent;
+  background-size: 200% auto;
+  background-clip: text;
+  -webkit-background-clip: text;
+  animation: glisten 3s linear infinite;
+}
+/* Overrides */
+.message h3[data-streamdown="heading-3"] {
+  @apply mt-3;
+}
+.message .collapsible p {
+  @apply my-1;
+}
+.message p {
+  @apply whitespace-pre-wrap;
+}
+@keyframes wobble {
+  0%,
+  100% {
+    transform: rotate(-3deg);
+  }
+  50% {
+    transform: rotate(3deg);
+  }
+}
+.animate-wobble {
+  animation: wobble 2s ease-in-out infinite;
+}

src/main.tsx ADDED Viewed

	@@ -0,0 +1,10 @@

+import { StrictMode } from "react";
+import { createRoot } from "react-dom/client";
+import "./index.css";
+import App from "./App.tsx";
+createRoot(document.getElementById("root")!).render(
+  <StrictMode>
+    <App />
+  </StrictMode>,
+);

tsconfig.app.json ADDED Viewed

	@@ -0,0 +1,28 @@

+{
+  "compilerOptions": {
+    "tsBuildInfoFile": "./node_modules/.tmp/tsconfig.app.tsbuildinfo",
+    "target": "ES2022",
+    "useDefineForClassFields": true,
+    "lib": ["ES2022", "DOM", "DOM.Iterable"],
+    "module": "ESNext",
+    "types": ["vite/client"],
+    "skipLibCheck": true,
+    /* Bundler mode */
+    "moduleResolution": "bundler",
+    "allowImportingTsExtensions": true,
+    "verbatimModuleSyntax": true,
+    "moduleDetection": "force",
+    "noEmit": true,
+    "jsx": "react-jsx",
+    /* Linting */
+    "strict": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "erasableSyntaxOnly": true,
+    "noFallthroughCasesInSwitch": true,
+    "noUncheckedSideEffectImports": true
+  },
+  "include": ["src"]
+}

tsconfig.json ADDED Viewed

	@@ -0,0 +1,7 @@

+{
+  "files": [],
+  "references": [
+    { "path": "./tsconfig.app.json" },
+    { "path": "./tsconfig.node.json" }
+  ]
+}

tsconfig.node.json ADDED Viewed

	@@ -0,0 +1,26 @@

+{
+  "compilerOptions": {
+    "tsBuildInfoFile": "./node_modules/.tmp/tsconfig.node.tsbuildinfo",
+    "target": "ES2023",
+    "lib": ["ES2023"],
+    "module": "ESNext",
+    "types": ["node"],
+    "skipLibCheck": true,
+    /* Bundler mode */
+    "moduleResolution": "bundler",
+    "allowImportingTsExtensions": true,
+    "verbatimModuleSyntax": true,
+    "moduleDetection": "force",
+    "noEmit": true,
+    /* Linting */
+    "strict": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "erasableSyntaxOnly": true,
+    "noFallthroughCasesInSwitch": true,
+    "noUncheckedSideEffectImports": true
+  },
+  "include": ["vite.config.ts"]
+}

vite.config.ts ADDED Viewed

	@@ -0,0 +1,8 @@

+import { defineConfig } from "vite";
+import react from "@vitejs/plugin-react";
+import tailwindcss from "@tailwindcss/vite";
+// https://vite.dev/config/
+export default defineConfig({
+  plugins: [react(), tailwindcss()],
+});