Merge pull request #67 from addyosmani/update-models

feat: update webllm package + added Llama3.1 & Qwen2 models
addyosmani · Jul 25, 2024 · d00e4a1 · d00e4a1
2 parents b19d081 + abc593c
commit d00e4a1
Show file tree

Hide file tree

Showing 8 changed files with 260 additions and 14 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -10,14 +10,15 @@
   },
   "dependencies": {
     "@hookform/resolvers": "^3.3.4",
-    "@mlc-ai/web-llm": "^0.2.42",
+    "@mlc-ai/web-llm": "^0.2.51",
     "@next/third-parties": "^14.2.3",
     "@radix-ui/react-avatar": "^1.0.4",
     "@radix-ui/react-dialog": "^1.0.5",
     "@radix-ui/react-dropdown-menu": "^2.0.6",
     "@radix-ui/react-icons": "^1.3.0",
     "@radix-ui/react-label": "^2.0.2",
     "@radix-ui/react-popover": "^1.0.7",
+    "@radix-ui/react-scroll-area": "^1.1.0",
     "@radix-ui/react-select": "^2.0.0",
     "@radix-ui/react-slot": "^1.0.2",
     "@radix-ui/react-switch": "^1.0.3",

diff --git a/src/components/chat/chat-topbar.tsx b/src/components/chat/chat-topbar.tsx
@@ -14,6 +14,7 @@ import { Sidebar } from "../sidebar";
 import { Message } from "ai/react";
 import useChatStore from "@/hooks/useChatStore";
 import { Models, Model } from "@/lib/models";
+import { Badge } from "../ui/badge";
 
 interface ChatTopbarProps {
   chatId?: string;
@@ -48,22 +49,32 @@ export default function ChatTopbar({ chatId, stop }: ChatTopbarProps) {
             aria-expanded={open}
             className="w-[180px] md:w-[300px] justify-between bg-accent dark:bg-card"
           >
-            <p className="truncate">{selectedModel.displayName}</p>
+            <div className="flex gap-2 items-center">
+              <p className="truncate">{selectedModel.displayName}</p>
+              {selectedModel.badge && (
+                <Badge>{selectedModel.badge}</Badge>
+              )}
+            </div>
             <CaretSortIcon className="ml-2 h-4 w-4 shrink-0 opacity-50" />
           </Button>
         </PopoverTrigger>
-        <PopoverContent className="w-[200px] lg:w-[300px] p-1">
+        <PopoverContent className="w-[200px] lg:w-[300px] max-h-96 overflow-y-scroll p-1">
           {Models.map((model) => (
             <Button
               key={model.name}
               variant="ghost"
-              className="w-full justify-start"
+              className="w-full justify-start flex gap-2 items-center"
               onClick={() => {
                 setSelectedModel(model);
                 setOpen(false);
               }}
             >
               {model.displayName}
+              {model.badge && (
+                <Badge>
+                  {model.badge}
+                </Badge>
+              )}
             </Button>
           ))}
         </PopoverContent>

diff --git a/src/components/ui/badge.tsx b/src/components/ui/badge.tsx
@@ -0,0 +1,36 @@
+import * as React from "react"
+import { cva, type VariantProps } from "class-variance-authority"
+
+import { cn } from "@/lib/utils"
+
+const badgeVariants = cva(
+  "inline-flex items-center rounded-md border px-2.5 py-0.5 text-xs font-semibold transition-colors focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2",
+  {
+    variants: {
+      variant: {
+        default:
+          "border-transparent bg-primary text-primary-foreground shadow hover:bg-primary/80",
+        secondary:
+          "border-transparent bg-secondary text-secondary-foreground hover:bg-secondary/80",
+        destructive:
+          "border-transparent bg-destructive text-destructive-foreground shadow hover:bg-destructive/80",
+        outline: "text-foreground",
+      },
+    },
+    defaultVariants: {
+      variant: "default",
+    },
+  }
+)
+
+export interface BadgeProps
+  extends React.HTMLAttributes<HTMLDivElement>,
+  VariantProps<typeof badgeVariants> { }
+
+function Badge({ className, variant, ...props }: BadgeProps) {
+  return (
+    <div className={cn(badgeVariants({ variant }), className)} {...props} />
+  )
+}
+
+export { Badge, badgeVariants }
diff --git a/src/hooks/useChatStore.ts b/src/hooks/useChatStore.ts
@@ -40,7 +40,7 @@ interface Actions {
 const useChatStore = create<State & Actions>()(
   persist(
     (set) => ({
-      selectedModel: Models[2],
+      selectedModel: Models[13],
       setSelectedModel: (model: Model) =>
         set((state: State) => ({
           selectedModel:

diff --git a/src/lib/models.ts b/src/lib/models.ts
@@ -1,6 +1,7 @@
 export interface Model {
   name: string;
   displayName: string;
+  badge?: string;
 }
 
 // source is AppConfig
@@ -22,6 +23,16 @@ export const Models: Model[] = [
     name: "RedPajama-INCITE-Chat-3B-v1-q4f32_1-MLC",
     displayName: "RedPajama 3B",
   },
+  {
+    name: "Qwen2-1.5B-Instruct-q4f32_1-MLC",
+    displayName: "Qwen2 1.5B Instruct",
+    badge: "New",
+  },
+  {
+    name: "Qwen2-7B-Instruct-q4f32_1-MLC",
+    displayName: "Qwen2 7B Instruct",
+    badge: "New",
+  },
   {
     name: "Mistral-7B-Instruct-v0.3-q4f16_1-MLC",
     displayName: "Mistral 7B Instruct",
@@ -50,4 +61,19 @@ export const Models: Model[] = [
     name: "Llama-3-70B-Instruct-q3f16_1-MLC",
     displayName: "Llama3 70B Instruct",
   },
+  {
+    name: "Llama-3.1-8B-Instruct-q4f32_1-MLC",
+    displayName: "Llama3.1 8B Instruct",
+    badge: "New",
+  },
+  {
+    name: "Llama-3.1-8B-Instruct-q4f32_1-MLC-1k",
+    displayName: "Llama3.1 8B Instruct - (1k)",
+    badge: "New",
+  },
+  {
+    name: "Llama-3.1-70B-Instruct-q3f16_1-MLC",
+    displayName: "Llama3.1 70B Instruct",
+    badge: "New",
+  }
 ];
diff --git a/src/lib/web-llm-helper.ts b/src/lib/web-llm-helper.ts
@@ -82,14 +82,13 @@ export default class WebLLMHelper {
           content:
             customizedInstructions && isCustomizedInstructionsEnabled
               ? "You are a helpful assistant. Assist the user with their questions. You are also provided with the following information from the user, keep them in mind for your responses: " +
-                customizedInstructions
+              customizedInstructions
               : "You are a helpful assistant. Assist the user with their questions.",
         },
         ...storedMessages,
         { role: "user", content: input },
       ],
       temperature: 0.6,
-      max_gen_len: 1024,
     });
     for await (const chunk of completion) {
       const delta = chunk.choices[0].delta.content;

diff --git a/src/lib/worker.ts b/src/lib/worker.ts
@@ -1,9 +1,9 @@
 // https://www.npmjs.com/package/@mlc-ai/web-llm
-import { MLCEngineWorkerHandler, MLCEngine } from "@mlc-ai/web-llm";
+import { WebWorkerMLCEngineHandler } from "@mlc-ai/web-llm";
+
 
 // Hookup an Engine to a worker handler
-const engine = new MLCEngine();
-const handler = new MLCEngineWorkerHandler(engine);
+const handler = new WebWorkerMLCEngineHandler();
 self.onmessage = (msg: MessageEvent) => {
   handler.onmessage(msg);
 };