From 7686c87840eddcd8fbfe8f6b2027f597f74d44a2 Mon Sep 17 00:00:00 2001
From: Jay Wang <jay@zijie.wang>
Date: Thu, 28 Mar 2024 21:16:12 -0400
Subject: [PATCH] Update webllm template format

Signed-off-by: Jay Wang <jay@zijie.wang>
---
 package.json        | 10 +++----
 src/llms/web-llm.ts | 66 ++++++++++++++++++++++++++++-----------------
 2 files changed, 47 insertions(+), 29 deletions(-)

diff --git a/package.json b/package.json
index a1fe23c..56d9e0e 100644
--- a/package.json
+++ b/package.json
@@ -16,7 +16,7 @@
   },
   "devDependencies": {
     "@floating-ui/dom": "^1.6.3",
-    "@mlc-ai/web-llm": "^0.2.26",
+    "@mlc-ai/web-llm": "^0.2.29",
     "@tiptap/core": "^2.2.4",
     "@tiptap/extension-bubble-menu": "^2.2.4",
     "@tiptap/extension-document": "^2.2.4",
@@ -34,7 +34,7 @@
     "@types/diff-match-patch": "^1.0.36",
     "@types/emoji-regex": "^9.2.0",
     "@types/uuid": "^9.0.8",
-    "@typescript-eslint/eslint-plugin": "^7.2.0",
+    "@typescript-eslint/eslint-plugin": "^7.4.0",
     "@webgpu/types": "^0.1.40",
     "@xiaohk/utils": "^0.0.6",
     "d3-array": "^3.2.4",
@@ -53,10 +53,10 @@
     "idb-keyval": "^6.2.1",
     "prettier": "^3.2.5",
     "tippy.js": "^6.3.7",
-    "typescript": "^5.4.2",
+    "typescript": "^5.4.3",
     "uuid": "^9.0.1",
-    "vite": "^5.1.6",
-    "vite-plugin-dts": "^3.7.3",
+    "vite": "^5.2.6",
+    "vite-plugin-dts": "^3.8.1",
     "vite-plugin-web-components-hmr": "^0.1.3"
   }
 }
diff --git a/src/llms/web-llm.ts b/src/llms/web-llm.ts
index 3e66ea2..32c766b 100644
--- a/src/llms/web-llm.ts
+++ b/src/llms/web-llm.ts
@@ -83,49 +83,67 @@ const APP_CONFIGS: webllm.AppConfig = {
   ]
 };
 
+enum Role {
+  user = 'user',
+  assistant = 'assistant'
+}
+
 const CONV_TEMPLATES: Record<
   SupportedLocalModel,
   Partial<ConvTemplateConfig>
 > = {
   [SupportedLocalModel['tinyllama-1.1b']]: {
-    system: '<|im_start|><|im_end|> ',
-    roles: ['<|im_start|>user', '<|im_start|>assistant'],
+    system_template: '<|im_start|><|im_end|> ',
+    roles: {
+      [Role.user]: '<|im_start|>user',
+      [Role.assistant]: '<|im_start|>assistant'
+    },
     offset: 0,
     seps: ['', ''],
-    separator_style: 'Two',
-    stop_str: '<|im_end|>',
-    add_bos: false,
-    stop_tokens: [2]
+    stop_str: ['<|im_end|>'],
+    stop_token_ids: [2]
   },
   [SupportedLocalModel['llama-2-7b']]: {
-    system: '[INST] <<SYS>><</SYS>>\n\n ',
-    roles: ['[INST]', '[/INST]'],
+    system_template: '[INST] <<SYS>><</SYS>>\n\n ',
+    roles: {
+      [Role.user]: '[INST]',
+      [Role.assistant]: '[/INST]'
+    },
     offset: 0,
     seps: [' ', ' '],
-    separator_style: 'Two',
-    stop_str: '[INST]',
-    add_bos: true,
-    stop_tokens: [2]
+    role_content_sep: ' ',
+    role_empty_sep: ' ',
+    stop_str: ['[INST]'],
+    system_prefix_token_ids: [1],
+    stop_token_ids: [2],
+    add_role_after_system_message: false
   },
   [SupportedLocalModel['phi-2']]: {
-    system: '',
-    roles: ['Instruct', 'Output'],
+    system_template: '',
+    system_message: '',
+    roles: {
+      [Role.user]: 'Instruct',
+      [Role.assistant]: 'Output'
+    },
     offset: 0,
     seps: ['\n'],
-    separator_style: 'Two',
-    stop_str: '<|endoftext|>',
-    add_bos: false,
-    stop_tokens: [50256]
+    stop_str: ['<|endoftext|>'],
+    stop_token_ids: [50256]
   },
   [SupportedLocalModel['gemma-2b']]: {
-    system: '',
-    roles: ['<start_of_turn>user', '<start_of_turn>model'],
+    system_template: '',
+    system_message: '',
+    roles: {
+      [Role.user]: '<start_of_turn>user',
+      [Role.assistant]: '<start_of_turn>model'
+    },
     offset: 0,
     seps: ['<end_of_turn>\n', '<end_of_turn>\n'],
-    separator_style: 'Two',
-    stop_str: '<end_of_turn>',
-    add_bos: true,
-    stop_tokens: [1, 107]
+    role_content_sep: '\n',
+    role_empty_sep: '\n',
+    stop_str: ['<end_of_turn>'],
+    system_prefix_token_ids: [2],
+    stop_token_ids: [1, 107]
   }
 };