fix wizard chat models

smallcloudai · Oct 14, 2023 · d7696ca · d7696ca
1 parent ad93630
commit d7696ca
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/src/caps.rs b/src/caps.rs
@@ -160,12 +160,13 @@ const KNOWN_MODELS: &str = r#"
             "n_ctx": 2048,
             "supports_scratchpads": {
                 "CHAT-GENERIC": {
-                    "token_esc": "\n\n",
-                    "keyword_system": "",
-                    "keyword_user": "USER: ",
-                    "keyword_assistant": "ASSISTANT: ",
-                    "stop_list": [],
-                    "default_system_message": ""
+                    "token_esc": "",
+                    "keyword_system": "<s>",
+                    "keyword_user": "\nUSER: ",
+                    "keyword_assistant": "\nASSISTANT: ",
+                    "eot": "",
+                    "stop_list": ["\n\n"],
+                    "default_system_message": "You are a helpful AI assistant.\n"
                 }
             },
             "similar_models": [

diff --git a/src/scratchpads/chat_generic.rs b/src/scratchpads/chat_generic.rs
@@ -55,9 +55,12 @@ impl ScratchpadAbstract for GenericChatScratchpad {
         self.keyword_asst = patch.get("keyword_assistant").and_then(|x| x.as_str()).unwrap_or("ASSISTANT:").to_string();
         self.default_system_message = patch.get("default_system_message").and_then(|x| x.as_str()).unwrap_or("").to_string();
         self.t.eot = patch.get("eot").and_then(|x| x.as_str()).unwrap_or("<|endoftext|>").to_string();
-        self.t.assert_one_token(&self.t.eot.as_str())?;
+
         self.dd.stop_list.clear();
-        self.dd.stop_list.push(self.t.eot.clone());
+        if !self.t.eot.is_empty() {
+            self.t.assert_one_token(&self.t.eot.as_str())?;
+            self.dd.stop_list.push(self.t.eot.clone());
+        }
         if self.token_esc.len() > 0 {
             self.dd.stop_list.push(self.token_esc.clone());
         } else {