show tokens per second

cfahlgren1 · Oct 18, 2024 · 57ee402 · 57ee402
1 parent 57dc198
commit 57ee402
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 16 deletions.
diff --git a/app/components/MarkdownRenderer.tsx b/app/components/MarkdownRenderer.tsx
@@ -19,17 +19,21 @@ interface CodeBlockProps {
   language: string;
   value: string;
 }
-
 const CodeBlock: React.FC<CodeBlockProps> = ({ language, value }) => {
   return (
-    <SyntaxHighlighter 
-      language={language} 
-      style={oneDark} 
-      showLineNumbers={true}
-      className="rounded text-xs overflow-x-auto"
-    >
-      {value}
-    </SyntaxHighlighter>
+    <div className="relative group">
+      <div className="absolute top-0 right-0 bg-emerald-700 text-emerald-50 rounded-bl px-1.5 py-0.5 text-[10px] font-mono tracking-wide opacity-70 transition-opacity group-hover:opacity-100">
+        {language}
+      </div>
+      <SyntaxHighlighter 
+        language={language} 
+        style={oneDark} 
+        showLineNumbers={true}
+        className="rounded-md text-xs overflow-x-auto pt-5"
+      >
+        {value}
+      </SyntaxHighlighter>
+    </div>
   );
 };
 
@@ -52,15 +56,15 @@ const MarkdownRenderer: React.FC<MarkdownRendererProps> = ({ content, className
           {...props}
         />
       ) : (
-        <code className="px-1 py-0.5 rounded bg-emerald-800 text-emerald-100 border border-emerald-600" {...props}>
+        <code className="px-1 py-0.5 rounded bg-emerald-700 text-emerald-100 text-xs" {...props}>
           {children}
         </code>
       );
     },
   };
 
   return (
-    <div className={`prose prose-invert max-w-none ${className || ''}`}>
+    <div className={`prose prose-sm prose-invert max-w-none ${className || ''}`}>
       <ReactMarkdown
         components={components}
         remarkPlugins={[remarkMath]}

diff --git a/app/globals.css b/app/globals.css
@@ -76,3 +76,13 @@ body {
     @apply bg-background text-foreground;
   }
 }
+
+::selection {
+  background-color: #34d399; /* emerald-400 */
+  color: #064e3b; /* emerald-900 */
+}
+
+::-moz-selection {
+  background-color: #34d399; /* emerald-400 */
+  color: #064e3b; /* emerald-900 */
+}
diff --git a/app/page.tsx b/app/page.tsx
@@ -10,9 +10,10 @@ import { Loader2 } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { handleAIStreaming } from "../lib/llm";
 import MarkdownRenderer from './components/MarkdownRenderer';
+import { CompletionUsage } from "@mlc-ai/web-llm";
 
 export default function Home() {
-  const [userInput, setUserInput] = useState<string>("");
+  const [userInput, setUserInput] = useState('');
   const [result, setResult] = useState<string | null>(null);
   const [resultExplanation, setResultExplanation] = useState<string | null>(null);
   const [hasError, setHasError] = useState(false);
@@ -22,6 +23,7 @@ export default function Home() {
   const [loadingProgress, setLoadingProgress] = useState<string>("");
   const [isStreaming, setIsStreaming] = useState(false);
   const [codeOutput, setCodeOutput] = useState<string | null>(null);
+  const [usage, setUsage] = useState<CompletionUsage | null>(null);
 
   useEffect(() => {
     setProgressCallback(setLoadingProgress);
@@ -56,6 +58,7 @@ export default function Home() {
     setResultExplanation("");
     setCodeOutput(null);
     setHasError(false);
+    setUsage(null);
 
     try {
       await handleAIStreaming(
@@ -66,6 +69,7 @@ export default function Home() {
           onCodeOutputUpdate: setCodeOutput,
           onExplanationUpdate: setResultExplanation,
           onErrorUpdate: setHasError,
+          onUsageUpdate: setUsage,
         }
       );
     } catch (err) {
@@ -83,6 +87,10 @@ export default function Home() {
     }
   };
 
+  const formatTokensPerSecond = (value: number) => {
+    return value.toFixed(2);
+  };
+
   return (
     <div className="flex flex-col min-h-screen bg-emerald-900">
       <main className="flex-grow flex flex-col items-center justify-center px-4 py-24">
@@ -146,12 +154,19 @@ export default function Home() {
         {codeOutput && (
           <div className="bg-teal-800 p-4 w-full max-w-2xl overflow-auto border-t border-teal-700">
             <h3 className="text-teal-200 font-semibold mb-2">Code Output:</h3>
-            <MarkdownRenderer content={`\`\`\`bash\n${codeOutput}\n\`\`\``} />
+            <MarkdownRenderer content={`\`\`\`output\n${codeOutput}\n\`\`\``} />
           </div>
         )}
         {resultExplanation && (
           <div className="bg-emerald-700 p-4 rounded-b w-full max-w-2xl overflow-auto border-t border-emerald-600">
-            <h3 className="text-emerald-200 font-semibold mb-2">Explanation:</h3>
+            <div className="flex justify-between items-center mb-2">
+              <h3 className="text-emerald-200 font-semibold">Explanation:</h3>
+              {usage && (
+                <div className="bg-emerald-600 text-emerald-50 rounded px-1.5 py-0.5 text-xs font-mono tracking-wide">
+                  {formatTokensPerSecond(usage.extra.decode_tokens_per_s)} tok/s
+                </div>
+              )}
+            </div>
             <MarkdownRenderer content={resultExplanation} className="text-emerald-50" />
           </div>
         )}

diff --git a/lib/llm.ts b/lib/llm.ts
@@ -93,6 +93,7 @@ export interface StreamingCallbacks {
   onCodeOutputUpdate: (output: string | null) => void;
   onExplanationUpdate: (explanation: string | null) => void;
   onErrorUpdate: (hasError: boolean) => void;
+  onUsageUpdate: (usage: webllm.CompletionUsage) => void;
 }
 
 export async function handleAIStreaming(
@@ -120,8 +121,9 @@ export async function handleAIStreaming(
       (currentMessage) => {
         callbacks.onResultUpdate(currentMessage);
       },
-      async (finalMessage: string, usage: unknown) => {
+      async (finalMessage: string, usage: webllm.CompletionUsage) => {
         callbacks.onResultUpdate(finalMessage);
+        callbacks.onUsageUpdate(usage);
         console.log("Usage:", usage);
 
         const extractedCode = extractCodeFromMarkdown(finalMessage);
@@ -148,8 +150,9 @@ export async function handleAIStreaming(
                 (currentExplanation) => {
                   callbacks.onExplanationUpdate(currentExplanation);
                 },
-                (finalExplanation: string) => {
+                (finalExplanation: string, explanationUsage: webllm.CompletionUsage) => {
                   callbacks.onExplanationUpdate(finalExplanation);
+                  callbacks.onUsageUpdate(explanationUsage);
                 },
                 (error: Error) => {
                   callbacks.onErrorUpdate(true);