added stats and template

.idea/vcs.xml

@@ -3,6 +3,5 @@
   <component name="VcsDirectoryMappings">
     <mapping directory="" vcs="Git" />
     <mapping directory="$PROJECT_DIR$" vcs="Git" />
-    <mapping directory="$PROJECT_DIR$/transformers.js-v4/transformers.js" vcs="Git" />
   </component>
 </project>

src/App.tsx

@@ -8,10 +8,10 @@ export default function App() {
   return (
     <ThemeContextProvider>
       <ChatSettingsContextProvider>
-        <div className="mx-auto flex h-screen w-full max-w-4xl flex-col">
-          <Header className="" />
-          <Chat className="flex-grow" />
+        <div className="fixed top-0 right-0 left-0 mx-auto w-full max-w-4xl">
+          <Header className="mx-auto w-full" />
         </div>
+        <Chat />
       </ChatSettingsContextProvider>
     </ThemeContextProvider>
   );

src/chat/Chat.tsx

@@ -1,7 +1,7 @@
 import { Button, Card, Slider } from "@theme";
 import { useChatSettings } from "@utils/context/chatSettings/useChatSettings.ts";
 import { formatBytes } from "@utils/format.ts";
-import { MODELS } from "@utils/models.ts";
+import { CONVERSATION_STARTERS, MODELS } from "@utils/models.ts";
 import { TOOLS } from "@utils/tools.ts";
 import { Settings } from "lucide-react";
 import {
@@ -13,10 +13,13 @@ import {
 } from "react";
 
 import TextGeneration from "../textGeneration/TextGeneration.ts";
+import type {
+  ChatMessageAssistant,
+  ChatMessageUser,
+} from "../textGeneration/types.ts";
 import cn from "../utils/classnames.ts";
 import ChatForm from "./ChatForm.tsx";
-import MessageContent from "./MessageContent.tsx";
-import MessageToolCall from "./MessageToolCall.tsx";
+import Message from "./Message.tsx";
 
 enum State {
   IDLE,
@@ -31,6 +34,8 @@ export default function Chat({ className = "" }: { className?: string }) {
   const [downloadProgress, setDownloadProgress] = useState<number>(0);
   const [state, setState] = useState<State>(State.IDLE);
   const settingsKey = useRef<string>(null);
+  const [starterCategory, setStarterCategory] = useState<number>(0);
+  const messagesEndRef = useRef<HTMLDivElement>(null);
 
   const generator = useMemo(() => new TextGeneration(), []);
   const getSnapshot = () => generator.chatMessages;
@@ -51,6 +56,10 @@ export default function Chat({ className = "" }: { className?: string }) {
     );
   }, [settings]);
 
+  useEffect(() => {
+    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
+  }, [messages]);
+
   if (!settings) return;
 
   const modelDownloaded = downloadedModels.includes(settings.modelKey);
@@ -78,11 +87,18 @@ export default function Chat({ className = "" }: { className?: string }) {
   const model = MODELS[settings.modelKey];
   const ready: boolean = state === State.READY || modelDownloaded;
 
+  const conversationMessages = messages.filter(({ role }) => role !== "system");
+
   return (
-    <div className={cn(className, "flex flex-col gap-4 px-4 py-8")}>
-      {state === State.IDLE && !modelDownloaded ? (
-        <div className="flex h-full items-center justify-center">
-          <Card className="flex max-w-[500px] flex-col gap-4">
+    <>
+      <div
+        className={cn(
+          className,
+          "mx-auto min-h-screen w-full max-w-4xl px-4 pt-24 pb-36"
+        )}
+      >
+        {state === State.IDLE && !modelDownloaded ? (
+          <Card className="mt-26 flex max-w-[500px] flex-col gap-4 self-center justify-self-center">
             <p>
               You are about to load <b>{model.title}</b>.<br />
               Once downloaded, the model ({formatBytes(model.size)}) will be
@@ -98,26 +114,58 @@ export default function Chat({ className = "" }: { className?: string }) {
               Download Model ({formatBytes(model.size)})
             </Button>
           </Card>
-        </div>
-      ) : state === State.INITIALIZING ? (
-        <div className="flex h-full items-center justify-center">
-          <div className="flex h-full w-full max-w-[400px] flex-col items-center justify-center gap-2">
-            <p className="flex w-full items-center justify-between">
-              <span>
-                {modelDownloaded
-                  ? "initializing the model..."
-                  : "downloading the model..."}
-              </span>
-              <span>{downloadProgress.toFixed(2)}%</span>
-            </p>
-            <Slider width={Math.round(downloadProgress)} />
+        ) : state === State.INITIALIZING ? (
+          <div className="mt-46 flex h-full items-center justify-center">
+            <div className="flex h-full w-full max-w-[400px] flex-col items-center justify-center gap-2">
+              <p className="flex w-full items-center justify-between">
+                <span>
+                  {modelDownloaded
+                    ? "initializing the model..."
+                    : "downloading the model..."}
+                </span>
+                <span>{downloadProgress.toFixed(2)}%</span>
+              </p>
+              <Slider width={Math.round(downloadProgress)} />
+            </div>
           </div>
-        </div>
-      ) : (
-        <div className="flex flex-col gap-4">
-          {messages
-            .filter(({ role }) => role !== "system")
-            .map((message, index) => (
+        ) : conversationMessages.length === 0 ? (
+          <div className="mt-26 flex h-full items-center justify-center">
+            <div className="flex flex-col gap-8">
+              <h2 className="text-2xl font-bold">How can I help you?</h2>
+              <div className="flex gap-4">
+                {CONVERSATION_STARTERS.map((starter, index) => (
+                  <Button
+                    iconLeft={starter.icon}
+                    key={index}
+                    onClick={() => setStarterCategory(index)}
+                    variant={index === starterCategory ? "solid" : "outline"}
+                  >
+                    {starter.category}
+                  </Button>
+                ))}
+              </div>
+              <div className="gap- flex flex-col border-t border-gray-300 dark:border-gray-700">
+                {CONVERSATION_STARTERS[starterCategory].prompts.map(
+                  (prompt, index) => (
+                    <div className="border-b border-gray-300 py-1 dark:border-gray-700">
+                      <Button
+                        key={index}
+                        onClick={() => generate(prompt)}
+                        variant="ghost"
+                        color="mono"
+                        className="justify-start"
+                      >
+                        {prompt}
+                      </Button>
+                    </div>
+                  )
+                )}
+              </div>
+            </div>
+          </div>
+        ) : (
+          <div className="flex flex-col gap-4">
+            {conversationMessages.map((message, index) => (
               <div
                 className={cn("w-3/4", {
                   "self-end": message.role === "user",
@@ -125,45 +173,41 @@ export default function Chat({ className = "" }: { className?: string }) {
                 key={index}
               >
                 {message.role === "user" ? (
-                  <Card>
-                    <p>{message.content}</p>
+                  <Card className="mt-2 mb-12">
+                    <p>{(message as ChatMessageUser).content}</p>
                   </Card>
                 ) : message.role === "assistant" ? (
-                  <div className="space-y-1">
-                    {message.content.map((part) =>
-                      part.type === "tool" ? (
-                        <MessageToolCall tool={part} />
-                      ) : (
-                        <MessageContent content={part.content} />
-                      )
-                    )}
-                  </div>
+                  <Message message={message as ChatMessageAssistant} />
                 ) : null}
               </div>
             ))}
+            <div ref={messagesEndRef} />
+          </div>
+        )}
+      </div>
+      <div className="fixed right-0 bottom-6 left-0 mx-auto w-full max-w-4xl">
+        <Card className="mt-auto">
+          <ChatForm
+            disabled={!ready}
+            onSubmit={(prompt) => generate(prompt)}
+            isGenerating={state === State.GENERATING}
+            onAbort={() => generator.abort()}
+          />
+        </Card>
+        <div className="flex justify-between dark:border-gray-700">
+          <Button
+            iconLeft={<Settings />}
+            size="xs"
+            variant="ghost"
+            color="mono"
+            onClick={openSettingsModal}
+          >
+            {settings?.modelKey ? MODELS[settings.modelKey].title : ""} |{" "}
+            {settings?.tools.length}/{TOOLS.length} tool
+            {settings?.tools.length !== 1 ? "s" : ""} active
+          </Button>
         </div>
-      )}
-      <Card className="mt-auto">
-        <ChatForm
-          disabled={!ready}
-          onSubmit={(prompt) => generate(prompt)}
-          isGenerating={state === State.GENERATING}
-          onAbort={() => generator.abort()}
-        />
-      </Card>
-      <div className="flex justify-between dark:border-gray-700">
-        <Button
-          iconLeft={<Settings />}
-          size="xs"
-          variant="ghost"
-          color="mono"
-          onClick={openSettingsModal}
-        >
-          {settings?.modelKey ? MODELS[settings.modelKey].title : ""} |{" "}
-          {settings?.tools.length}/{TOOLS.length} tool
-          {settings?.tools.length !== 1 ? "s" : ""} active
-        </Button>
       </div>
-    </div>
+    </>
   );
 }

src/chat/Message.tsx

@@ -0,0 +1,75 @@
+import { Modal } from "@theme";
+import cn from "@utils/classnames.ts";
+import { SquareCode, StopCircle } from "lucide-react";
+import { useState } from "react";
+
+import type { ChatMessageAssistant } from "../textGeneration/types.ts";
+import MessageContent from "./MessageContent.tsx";
+import MessageToolCall from "./MessageToolCall.tsx";
+
+export default function Message({
+  message,
+  className = "",
+}: {
+  message: ChatMessageAssistant;
+  className?: string;
+}) {
+  const [templateOpen, setTemplateOpen] = useState<boolean>(false);
+  const formatDuration = (ms: number) => {
+    if (ms < 1000) {
+      return `${Math.round(ms)}ms`;
+    }
+    return `${(ms / 1000).toFixed(2)}s`;
+  };
+
+  return (
+    <div className={cn(className, "group relative flex flex-col gap-1")}>
+      <Modal
+        title="Full Template"
+        isOpen={templateOpen}
+        onClose={() => setTemplateOpen(false)}
+        size="xl"
+      >
+        <pre className="max-h-[70vh] overflow-auto rounded-lg border border-gray-300 bg-gray-50 p-4 text-sm leading-relaxed text-gray-900 dark:border-gray-600 dark:bg-gray-900 dark:text-gray-100">
+          {message.template}
+        </pre>
+      </Modal>
+      {message.content.map((part) =>
+        part.type === "tool" ? (
+          <MessageToolCall tool={part} />
+        ) : (
+          <MessageContent content={part.content} />
+        )
+      )}
+      {message.interrupted && (
+        <div className="mt-2 flex items-center gap-2 rounded-md border border-amber-300 bg-amber-50 px-3 py-2 text-sm text-amber-800 dark:border-amber-700 dark:bg-amber-950 dark:text-amber-200">
+          <StopCircle className="size-4 shrink-0" />
+          <span>Response interrupted by the user</span>
+        </div>
+      )}
+      {message.modelUsage && (
+        <div className="mt-1 inline-flex flex-wrap items-center gap-1.5 self-end rounded-lg bg-gray-200 px-2.5 py-1 text-xs font-medium text-gray-700 opacity-0 transition-opacity group-hover:opacity-100 dark:bg-gray-700 dark:text-gray-300">
+          <span className="font-semibold">{message.modelUsage.model}</span>
+          <span className="text-gray-400 dark:text-gray-500">•</span>
+          <span>
+            Output: {Math.round(message.modelUsage.outputTps)} tok/s (
+            {message.modelUsage.outputTokens} tokens in{" "}
+            {formatDuration(message.modelUsage.outputDurationMs)})
+          </span>
+          <span className="text-gray-400 dark:text-gray-500">•</span>
+          <span title="Input processing time">
+            Prefill: {formatDuration(message.modelUsage.inputDurationMs)}
+          </span>
+          <span className="text-gray-400 dark:text-gray-500">•</span>
+          <button
+            className="flex items-center gap-1 rounded-md px-1.5 py-0.5 transition-colors hover:bg-gray-300 hover:text-gray-900 dark:hover:bg-gray-600 dark:hover:text-gray-100"
+            onClick={() => setTemplateOpen(true)}
+          >
+            <SquareCode className="size-3 -translate-y-1/20" />
+            <span>Template</span>
+          </button>
+        </div>
+      )}
+    </div>
+  );
+}

src/chat/MessageContent.tsx

@@ -95,12 +95,15 @@ export default function MessageContent({ content }: { content: string }) {
           className="-ml-2 opacity-50"
           loading={isThinking}
           iconLeft={<Lightbulb />}
+          notDisabledWhileLoading
         >
           {isThinking ? "Thinking for" : "Thought for"}{" "}
           {thinkingTime.toFixed(1)}s ({showThinking ? "Hide" : "Show"})
         </Button>
         {showThinking && (
-          <p className="max-w-none text-sm">{parsed.thinkContent}</p>
+          <p className="my-4 max-w-none border-l-1 border-gray-300 pl-4 text-sm font-medium text-gray-600 dark:border-gray-700 dark:text-gray-400">
+            {parsed.thinkContent}
+          </p>
         )}
       </div>
       {parsed.afterContent && (

src/index.css

@@ -23,7 +23,7 @@
 }
 
 body {
-  background-color: white;
+  background-color: #f6f6f6;
   color: black;
   --shiny-width: 50px;
 }

src/textGeneration/TextGeneration.ts

@@ -145,6 +145,7 @@ export default class TextGeneration {
       response: string;
       modelUsage: ModelUsage;
       interrupted: boolean;
+      template: string;
     }>((resolve, reject) => {
       if (this.modelKey === null) {
         reject("Model not initialized");
@@ -162,7 +163,6 @@ export default class TextGeneration {
 
       const requestId = this.requestId++;
       this.messages = [...this.messages, { role, content: prompt }];
-      this.messages.push({ role: "assistant", content: "" });
 
       let response = "";
 
@@ -174,10 +174,12 @@ export default class TextGeneration {
         }
         if (data.type === ResponseType.GENERATE_TEXT_DONE) {
           this.removeWorkerEventListener(listener);
+          this.messages.push({ role: "assistant", content: data.response });
           resolve({
             response: data.response,
             modelUsage: data.modelUsage,
             interrupted: data.interrupted,
+            template: data.template,
           });
         }
         if (data.type === ResponseType.GENERATE_TEXT_CHUNK) {
@@ -217,7 +219,7 @@ export default class TextGeneration {
 
       this.chatMessages = [...prevChatMessages, assistantMessage];
 
-      const { modelUsage, interrupted } = await this.generateText(
+      const { modelUsage, interrupted, template } = await this.generateText(
         prompt,
         isUser ? "user" : "tool",
         (partialResponse) => {
@@ -242,6 +244,7 @@ export default class TextGeneration {
       );
       isUser = false;
 
+      assistantMessage.template = template;
       assistantMessage.modelUsage = modelUsage;
       assistantMessage.interrupted = interrupted;
       this.chatMessages = [...prevChatMessages, assistantMessage];

src/textGeneration/types.ts

@@ -59,6 +59,7 @@ interface ResponseGenerateTextDone {
   type: ResponseType.GENERATE_TEXT_DONE;
   response: string;
   modelUsage: ModelUsage;
+  template: string;
   interrupted: boolean;
   requestId: number;
 }
@@ -107,6 +108,7 @@ export interface ChatMessageAssistant {
   content: Array<ChatMessageAssistantResponse | ChatMessageAssistantTool>;
   interrupted: boolean;
   modelUsage?: ModelUsage;
+  template?: string;
 }
 
 export interface ChatMessageSystem {

src/textGeneration/worker/textGenerationWorker.ts

@@ -7,9 +7,9 @@ import {
   Tensor,
   TextStreamer,
 } from "@huggingface/transformers";
+
 import { calculateDownloadProgress } from "../../utils/calculateDownloadProgress.ts";
 import { MODELS } from "../../utils/models.ts";
-
 import {
   type Request,
   RequestType,
@@ -100,7 +100,6 @@ self.onmessage = async ({ data }: MessageEvent<Request>) => {
     stoppingCriteria = new InterruptableStoppingCriteria();
 
     const { messages, tools, requestId } = data;
-    console.log(messages, tools);
     const { tokenizer, model } = await getTextGenerationPipeline(data.modelKey);
     if (!stoppingCriteria) {
       stoppingCriteria = new InterruptableStoppingCriteria();
@@ -170,6 +169,10 @@ self.onmessage = async ({ data }: MessageEvent<Request>) => {
       }
     )[0];
 
+    const template = tokenizer.batch_decode(sequences, {
+      skip_special_tokens: false,
+    })[0];
+
     cache = {
       pastKeyValues: past_key_values,
       key:
@@ -183,20 +186,19 @@ self.onmessage = async ({ data }: MessageEvent<Request>) => {
         ]),
     };
 
-    console.log("response", response);
-
     postMessage({
       type: ResponseType.GENERATE_TEXT_DONE,
       response,
       modelUsage: {
-        inputDurationMs: (firstTokenTime || 0) - started,
+        inputDurationMs: firstTokenTime - started,
         outputTokens: numTokens,
-        outputDurationMs: performance.now() - ended,
+        outputDurationMs: ended - firstTokenTime,
         outputTps: tps,
         doneMs: ended - started,
         modelKey: MODEL.modelId,
         model: MODEL.title,
       },
+      template,
       interrupted: stoppingCriteria.interrupted,
       requestId,
     });

src/theme/button/Button.tsx

@@ -17,6 +17,7 @@ interface BaseButtonProps {
   iconRight?: ReactNode;
   disabled?: boolean;
   loading?: boolean;
+  notDisabledWhileLoading?: boolean;
   shiny?: boolean;
 }
 
@@ -104,6 +105,7 @@ export default function Button({
   iconRight,
   disabled = false,
   loading = false,
+  notDisabledWhileLoading = false,
   shiny = false,
   ...props
 }: ButtonProps) {
@@ -119,6 +121,8 @@ export default function Button({
     className
   );
 
+  const isDisabled = disabled || (loading && !notDisabledWhileLoading);
+
   const renderIcon = (icon: ReactNode) => {
     if (isValidElement(icon)) {
       return cloneElement(icon as any, {
@@ -144,7 +148,7 @@ export default function Button({
         className={baseClasses}
         target={props.target}
         rel={props.rel}
-        aria-disabled={disabled || loading}
+        aria-disabled={isDisabled}
       >
         {content}
       </a>
@@ -155,7 +159,7 @@ export default function Button({
     <button
       type={(props as ButtonAsButton).type || "button"}
       onClick={(props as ButtonAsButton).onClick}
-      disabled={disabled || loading}
+      disabled={isDisabled}
       className={baseClasses}
     >
       {content}

src/theme/misc/Card.tsx

@@ -12,7 +12,7 @@ export default function Card({
     <div
       className={cn(
         className,
-        "rounded-lg border border-gray-200 bg-white p-4 dark:border-gray-700 dark:bg-gray-800"
+        "rounded-lg border border-gray-300 bg-gray-100 p-4 shadow-md dark:border-gray-700 dark:bg-gray-800"
       )}
     >
       {children}

src/utils/models.ts

@@ -1,3 +1,6 @@
+import { CircleQuestionMark, Pen, Sparkles } from "lucide-react";
+import { createElement } from "react";
+
 export interface Model {
   modelId: string;
   title: string;
@@ -135,3 +138,36 @@ export const MODELS: Record<string, Model> = {
 export type ModelKey = keyof typeof MODELS;
 
 export const DEFAULT_SYSTEM_PROMPT = "You are a helpful AI Assistant";
+
+export const CONVERSATION_STARTERS = [
+  {
+    category: "Getting Started",
+    icon: createElement(Pen),
+    prompts: [
+      "Tell me something about yourself. Who are you and who trained you?",
+      "What can you help me with today?",
+      "Explain how you work in simple terms",
+      "What are your strengths and limitations?",
+    ],
+  },
+  {
+    category: "Creative Writing",
+    icon: createElement(Sparkles),
+    prompts: [
+      "Write a short story about a robot discovering emotions",
+      "Create a character profile for a friendly neighborhood baker",
+      "Give me 3 creative writing prompts for short stories",
+      "Write a dialogue between two friends planning a road trip",
+    ],
+  },
+  {
+    category: "Everyday Help",
+    icon: createElement(CircleQuestionMark),
+    prompts: [
+      "Help me write a polite email to reschedule a meeting",
+      "Give me 5 tips for staying productive while working from home",
+      "Suggest some healthy meal ideas for busy weekdays",
+      "How do I organize my daily tasks more effectively?",
+    ],
+  },
+];