chore: remove dup with_stop_words

TabbyML · Oct 4, 2024 · abcfe6b · abcfe6b
1 parent 7765b55
commit abcfe6b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 15 deletions.
diff --git a/crates/tabby-index/src/code/snapshots/tabby_index__code__index__tests__code_splitter.snap b/crates/tabby-index/src/code/snapshots/tabby_index__code__index__tests__code_splitter.snap
@@ -1,6 +1,5 @@
 ---
 source: crates/tabby-index/src/code/index.rs
-assertion_line: 161
 expression: "format!(\"{:#?}\", rust_chunks)"
 ---
 [
@@ -23,5 +22,5 @@ expression: "format!(\"{:#?}\", rust_chunks)"
     "{\n        let text = reverse(\"void write_u32(std::uint32_t val) const {\\n        write_raw(&val, sizeof(val));\\n    }\\n\\n    ~llama_file() {\\n        if (fp) {\\n            std::fclose(fp);\\n        }\\n    }\\n};\\n\\nvoid\");\n\n        let trie = create_stop_trie(vec![\"\\n\\n\".to_owned(), \"\\n\\n  \".to_owned()]);\n        assert!(trie.common_prefix_search(&text).is_empty());",
     "let trie = create_stop_trie(vec![\n            \"\\n\\n\".to_owned(),\n            \"\\n\\n  \".to_owned(),\n            \"\\nvoid\".to_owned(),\n            \"<|file_sep|>\".to_owned(), // qwen 2.5 coder style\n        ]);\n        assert!(!trie.common_prefix_search(&text).is_empty());\n\n        let qwen25coder = reverse(\"qwen25 style stop words;<|file_sep|>\");\n        assert!(!trie.common_prefix_search(&qwen25coder).is_empty());\n    }\n\n    #[test]",
     "fn test_stop_condition_max_length() {\n        let factory = StopConditionFactory::default();\n        let mut cond = factory.create(\"\", Some(&UNKNOWN_LANGUAGE));\n        let (should_stop, _) = cond.should_stop(\"1\");\n        assert!(!should_stop);\n        let (should_stop, _) = cond.should_stop(\"2\");\n        assert!(!should_stop);\n        let (should_stop, _) = cond.should_stop(\"3\");\n        assert!(!should_stop);\n        let (should_stop, _) = cond.should_stop(\"4\");\n        assert!(!should_stop)\n    }",
-    "#[test]\n    fn test_stop_condition_additional_stop_words() {\n        let factory = StopConditionFactory::with_stop_words(vec![\"<|endoftext|>\".to_owned()]);\n        let mut cond = factory.create(\"additional_stop_words\", Some(&UNKNOWN_LANGUAGE));\n        let (should_stop, _) = cond.should_stop(\"1\");\n        assert!(!should_stop);\n        let (should_stop, _) = cond.should_stop(\"<|endoftext|>\");\n        assert!(should_stop);\n    }\n}",
+    "#[test]\n    fn test_stop_condition_additional_stop_words() {\n        let factory = StopConditionFactory::with_stop_words(vec![\"<|endoftext|>\".to_owned()]);\n        let mut cond = factory.create(\"\", Some(&UNKNOWN_LANGUAGE));\n        let (should_stop, _) = cond.should_stop(\"1\");\n        assert!(!should_stop);\n        let (should_stop, _) = cond.should_stop(\"<|endoftext|>\");\n        assert!(should_stop);\n    }\n}",
 ]
diff --git a/crates/tabby-inference/src/code.rs b/crates/tabby-inference/src/code.rs
@@ -32,15 +32,12 @@ pub struct CodeGeneration {
 
 impl CodeGeneration {
     pub fn new(imp: Arc<dyn CompletionStream>, config: Option<ModelConfig>) -> Self {
-        let stop_condition_factory = match config {
-            Some(ModelConfig::Local(config)) => StopConditionFactory::with_stop_words(
-                config.additional_stop_words.unwrap_or_default(),
-            ),
-            Some(ModelConfig::Http(config)) => StopConditionFactory::with_stop_words(
-                config.additional_stop_words.unwrap_or_default(),
-            ),
-            _ => StopConditionFactory::default(),
+        let additional_stop_words = match config {
+            Some(ModelConfig::Local(config)) => config.additional_stop_words.unwrap_or_default(),
+            Some(ModelConfig::Http(config)) => config.additional_stop_words.unwrap_or_default(),
+            _ => vec![],
         };
+        let stop_condition_factory = StopConditionFactory::with_stop_words(additional_stop_words);
 
         Self {
             imp,

diff --git a/crates/tabby/src/services/model/mod.rs b/crates/tabby/src/services/model/mod.rs
@@ -18,13 +18,14 @@ pub async fn load_code_generation_and_chat(
     Option<PromptInfo>,
     Option<Arc<dyn ChatCompletionStream>>,
 ) {
-    let (engine, prompt_info, chat) = load_completion_and_chat(&completion_model, chat_model).await;
+    let (engine, prompt_info, chat) =
+        load_completion_and_chat(completion_model.clone(), chat_model).await;
     let code = engine.map(|engine| Arc::new(CodeGeneration::new(engine, completion_model)));
     (code, prompt_info, chat)
 }
 
 async fn load_completion_and_chat(
-    completion_model: &Option<ModelConfig>,
+    completion_model: Option<ModelConfig>,
     chat_model: Option<ModelConfig>,
 ) -> (
     Option<Arc<dyn CompletionStream>>,
@@ -42,9 +43,9 @@ async fn load_completion_and_chat(
     let (completion, prompt) = if let Some(completion_model) = completion_model {
         match completion_model {
             ModelConfig::Http(http) => {
-                let engine = http_api_bindings::create(http).await;
+                let engine = http_api_bindings::create(&http).await;
                 let (prompt_template, chat_template) =
-                    http_api_bindings::build_completion_prompt(http);
+                    http_api_bindings::build_completion_prompt(&http);
                 (
                     Some(engine),
                     Some(PromptInfo {
@@ -54,7 +55,7 @@ async fn load_completion_and_chat(
                 )
             }
             ModelConfig::Local(llama) => {
-                let (stream, prompt) = llama_cpp_server::create_completion(llama).await;
+                let (stream, prompt) = llama_cpp_server::create_completion(&llama).await;
                 (Some(stream), Some(prompt))
             }
         }