Improve loading (#6)

* Disable input when no model loaded (115e780) * Disable Shiki (f803f7f) * Fix mobile height (075a611) * Update README (4da1500) * 0.1.2 (962cb82)
adamelliotfields · Feb 21, 2024 · 660a862 · 660a862
1 parent f48529c
commit 660a862
Show file tree

Hide file tree

Showing 7 changed files with 94 additions and 38 deletions.
diff --git a/index.html b/index.html
@@ -1,5 +1,5 @@
 <!doctype html>
-<html lang="en" class="h-full antialiased [font-synthesis:none] [text-rendering:optimizeLegibility]">
+<html lang="en" class="antialiased [font-synthesis:none] [text-rendering:optimizeLegibility]" style="height: 100%;">
   <head>
     <title>Chat</title>
     <meta charset="utf-8" />
@@ -10,8 +10,23 @@
     <!-- <link rel="icon" type="image/svg+xml" href="/favicon.svg" /> -->
     <link rel="icon" href="https://aef.me/favicon.ico" />
     <link rel="canonical" href="https://aef.me/chat/" />
-  <body class="h-full bg-neutral-50">
-    <div id="root" class="h-full"></div>
-    <script type="module" src="./src/index.tsx"></script>
+  <body class="bg-neutral-50" style="height: 100%;">
+    <div id="root" style="height: 100%;">
+      <span style="/* neutral-600 */ color: #525252; font-size: 32px; margin: 16px; display: block;">
+        <!-- https://github.com/cyberalien/line-md/blob/master/svg/loading-twotone-loop.svg -->
+        <svg xmlns="http://www.w3.org/2000/svg" width="1em" height="1em" viewBox="0 0 24 24">
+          <g fill="none" stroke="currentColor" stroke-linecap="round" stroke-width="2">
+            <path stroke-dasharray="60" stroke-dashoffset="60" stroke-opacity="0.3" d="M12 3C16.9706 3 21 7.02944 21 12C21 16.9706 16.9706 21 12 21C7.02944 21 3 16.9706 3 12C3 7.02944 7.02944 3 12 3Z">
+              <animate fill="freeze" attributeName="stroke-dashoffset" dur="1.3s" values="60;0" />
+            </path>
+            <path stroke-dasharray="15" stroke-dashoffset="15" d="M12 3C16.9706 3 21 7.02944 21 12">
+              <animate fill="freeze" attributeName="stroke-dashoffset" dur="0.3s" values="15;0" />
+              <animateTransform attributeName="transform" dur="1.5s" repeatCount="indefinite" type="rotate" values="0 12 12;360 12 12" />
+            </path>
+          </g>
+        </svg>
+      </span>
+    </div>
+    <script src="./src/index.tsx" type="module" async></script>
   </body>
 </html>
diff --git a/package.json b/package.json
@@ -1,7 +1,7 @@
 {
   "private": true,
   "name": "chat",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "type": "module",
   "scripts": {
     "start": "vite",

diff --git a/readme.md b/readme.md
@@ -2,7 +2,7 @@
 
 [![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/adamelliotfields/chat?devcontainer_path=.devcontainer/devcontainer.json&machine=basicLinux32gb)
 
-Static chat UI for [Web LLM](https://webllm.mlc.ai) on GitHub Pages. Inspired by [Perplexity Labs](https://labs.perplexity.ai).
+React chat UI for [Web LLM](https://webllm.mlc.ai) on GitHub Pages. Inspired by [Perplexity Labs](https://labs.perplexity.ai).
 
 https://github.com/adamelliotfields/chat/assets/7433025/07565763-606b-4de3-aa2d-8d5a26c83941
 
@@ -129,11 +129,8 @@ See [utils/vram_requirements](https://github.com/mlc-ai/web-llm/tree/main/utils/
 
 - [ ] Dark mode
 - [ ] Settings menu (temperature, system message, etc.)
-- [ ] Adapters for alternative backends (e.g., Ollama)
 - [ ] Inference on web worker
 - [ ] Offline/PWA
 - [ ] Cache management
-- [ ] GPU stats
 - [ ] Image upload for multimodal like [LLaVA](https://llava-vl.github.io)
-- [ ] [StableLM Zephyr 3B](https://huggingface.co/stabilityai/stablelm-zephyr-3b)
 - [ ] Tailwind class sorting by Biome 🤞
diff --git a/src/components/App.tsx b/src/components/App.tsx
@@ -2,7 +2,7 @@ import { ChatModule } from '@mlc-ai/web-llm'
 import { useAtom, useAtomValue, useSetAtom } from 'jotai'
 import { RESET } from 'jotai/utils'
 import { Square, Trash } from 'lucide-react'
-import { type MouseEvent } from 'react'
+import { type MouseEvent, useEffect } from 'react'
 
 import {
   activeModelIdAtom,
@@ -32,13 +32,28 @@ export default function App({ chat }: AppProps) {
   const [generating, setGenerating] = useAtom(generatingAtom)
   const [loading, setLoading] = useAtom(loadingAtom)
   const [conversation, setConversation] = useAtom(conversationAtom)
-  const setActiveModelId = useSetAtom(activeModelIdAtom)
+  const [activeModelId, setActiveModelId] = useAtom(activeModelIdAtom)
   const setStatsText = useSetAtom(runtimeStatsTextAtom)
   const config = useAtomValue(configAtom)
 
   const trashDisabled = conversation.messages.length < 1
   const stopDisabled = loading || !generating
 
+  // set initial status message
+  useEffect(() => {
+    const content =
+      "### Welcome\n\nThis app runs _small_ LLMs in your browser using your device's GPU. Select a model and press the power button to load it.\n\nErrors? Check [webgpureport.org](https://webgpureport.org) to inspect your system. A VPN can get around some network issues.\n\nRefresh the page to see this message again."
+    setConversation(() => ({
+      messages: [
+        {
+          messageRole: 'status',
+          content
+        }
+      ],
+      stream: null
+    }))
+  }, [setConversation])
+
   const onGenerate: GenerateCallback = (_, content) => {
     setConversation(({ messages }) => ({
       messages,
@@ -185,13 +200,13 @@ export default function App({ chat }: AppProps) {
           <div className="w-full max-w-screen-lg mx-auto">
             <MessageList
               // change `scroll` to `auto` if you don't like the scrollbars being always visible
-              className="h-[calc(100vh_-_56px_-_192px)] overflow-y-scroll md:h-[calc(100vh_-_56px_-_160px)]"
+              className="h-[calc(100vh_-_56px_-_204px)] overflow-y-scroll md:h-[calc(100vh_-_56px_-_160px)]"
             />
           </div>
         </main>
       </div>
 
-      <footer className="h-[192px] bg-neutral-50 sticky bottom-0 border-t z-20 md:h-[160px]">
+      <footer className="h-[204px] bg-neutral-50 sticky bottom-0 border-t z-20 md:h-[160px]">
         {/* footer top row */}
         <div className="max-w-screen-lg mx-auto">
           <div className="p-4 flex flex-col justify-between md:border-x md:flex-row md:items-center">
@@ -212,7 +227,10 @@ export default function App({ chat }: AppProps) {
                 onClick={handleResetClick}
               />
               <div className="grow">
-                <PromptInput handleClick={handleInputButtonClick} />
+                <PromptInput
+                  activeModelId={activeModelId}
+                  handleClick={handleInputButtonClick}
+                />
               </div>
               <Button
                 disabled={stopDisabled}

diff --git a/src/components/MessageList.tsx b/src/components/MessageList.tsx
@@ -1,4 +1,4 @@
-import Shiki from '@shikijs/markdown-it'
+// import Shiki from '@shikijs/markdown-it'
 import clsx from 'clsx'
 import { useAtomValue } from 'jotai'
 import MarkdownIt from 'markdown-it'
@@ -8,16 +8,33 @@ import { conversationAtom } from '../atoms'
 import type { Message as IMessage } from '../types'
 
 // can't use `react-markdown` with shiki or rehype-pretty until remarkjs/react-markdown#680
-// (markdown-it is awesome; I'm just depending on `dangerouslySetInnerHTML` to use it)
 const md = MarkdownIt()
-md.use(
-  await Shiki({
-    themes: {
-      light: 'github-light',
-      dark: 'github-dark'
-    }
-  })
-)
+
+// TODO: disabled for performance
+// md.use(
+//   await Shiki({
+//     themes: {
+//       light: 'github-light',
+//       dark: 'github-dark'
+//     }
+//   })
+// )
+
+// external links
+// https://github.com/markdown-it/markdown-it/blob/master/docs/architecture.md#renderer
+const renderToken =
+  md.renderer.rules.link_open ||
+  ((tokens, idx, options, _, self) => self.renderToken(tokens, idx, options))
+
+md.renderer.rules.link_open = (tokens, idx, options, env, self) => {
+  // attrs is a list of [name, value] tuples
+  const href = tokens[idx].attrs?.[0]?.[1]
+  if (href && !href.startsWith('#')) {
+    tokens[idx].attrPush(['target', '_blank'])
+    tokens[idx].attrPush(['rel', 'noopener noreferrer'])
+  }
+  return renderToken(tokens, idx, options, env, self)
+}
 
 export interface MessageProps extends IMessage {
   markdown?: boolean
@@ -34,15 +51,15 @@ export function Message({ content, markdown = false, messageRole }: MessageProps
       >
         {/* message bubble */}
         <div className="max-w-full border px-4 py-2 rounded-lg break-words [word-break:break-word] shadow-sm bg-white">
-          <div className="mt-1 text-xs text-neutral-400 font-bold tracking-tight uppercase">
+          <div className="mt-2 text-xs text-neutral-400 font-bold tracking-tight uppercase">
             {messageRole}
           </div>
           <div
             // biome-ignore lint: security/noDangerouslySetInnerHtml
             dangerouslySetInnerHTML={{
               __html: markdown ? md.render(content ?? '') : content ?? ''
             }}
-            className="mt-2 mb-0.5 prose min-w-0"
+            className="mt-4 mb-1 prose min-w-0"
           />
         </div>
       </div>

diff --git a/src/components/PromptInput.tsx b/src/components/PromptInput.tsx
@@ -5,20 +5,28 @@ import { type ChangeEvent, type HTMLAttributes, type MouseEvent, useState } from
 import { Button } from './Button'
 
 export interface PromptInputProps extends HTMLAttributes<HTMLDivElement> {
+  activeModelId?: string | null
   handleClick: (
     event: MouseEvent<HTMLButtonElement>,
     content: string
   ) => Promise<void> | void
 }
 
-export function PromptInput({ className, handleClick, ...rest }: PromptInputProps) {
+export function PromptInput({
+  activeModelId = null,
+  className,
+  handleClick,
+  ...rest
+}: PromptInputProps) {
   const [value, setValue] = useState<string>('')
-  const disabled = value.length < 1
+  const buttonDisabled = value.length < 1
+  const inputDisabled = activeModelId === null
+  const placeholder = inputDisabled ? 'Load a model...' : 'Ask anything...'
 
   const handleChange = (e: ChangeEvent<HTMLTextAreaElement>) => setValue(e.target.value)
 
   const handleClickWithValue = (e: MouseEvent<HTMLButtonElement>) => {
-    if (disabled) return
+    if (buttonDisabled) return
     handleClick(e, value)
     setValue('')
   }
@@ -35,21 +43,22 @@ export function PromptInput({ className, handleClick, ...rest }: PromptInputProp
       <textarea
         className={clsx(
           'h-10 flex-grow flex-shrink p-2 overflow-auto outline-none w-full font-sans resize-none caret-cyan-500 bg-white text-neutral-900 border-0',
-          'focus:ring-0',
-          'placeholder:text-neutral-400/50'
+          'focus:ring-0 placeholder:text-neutral-400/50',
+          inputDisabled && 'cursor-not-allowed'
         )}
         autoComplete="off"
-        placeholder="Ask anything..."
+        placeholder={placeholder}
         value={value}
         onChange={handleChange}
+        disabled={inputDisabled}
       />
       <div className="flex items-center justify-self-end bg-white rounded-full space-x-2">
         <Button
           className="text-lg"
-          disabled={disabled}
           label="Send"
           icon={SendHorizontal}
           onClick={handleClickWithValue}
+          disabled={buttonDisabled}
         />
       </div>
     </div>

diff --git a/src/components/RuntimeStats.tsx b/src/components/RuntimeStats.tsx
@@ -17,13 +17,13 @@ export function RuntimeStats({ className }: RuntimeStatsProps) {
     <div className={clsx('ml-1 flex items-center md:ml-5', className)}>
       {/* prefilling stats */}
       <div className={classNames}>
-        <span className="text-neutral-900">{stats.prefill} tokens/sec</span>
-        <span className="text-neutral-400">{' prefilling'}</span>
+        <span className="text-neutral-400 block lg:inline">prefilling</span>
+        <span className="text-neutral-900">{` ${stats.prefill}`} tokens/sec</span>
       </div>
       {/* decoding stats */}
-      <div className={clsx('ml-4 md:ml-0', classNames)}>
-        <span className="text-neutral-900">{stats.decode} tokens/sec</span>
-        <span className="text-neutral-400">{' decoding'}</span>
+      <div className={clsx('ml-8 md:ml-4', classNames)}>
+        <span className="text-neutral-400 block lg:inline">decoding</span>
+        <span className="text-neutral-900">{` ${stats.decode}`} tokens/sec</span>
       </div>
     </div>
   )