Merge pull request #457 from dcSpark/nico/add_gemma2_support

add gemma2 support
dcSpark · Jun 28, 2024 · 66e0ac4 · 66e0ac4
2 parents 24b08d0 + f10f9a8
commit 66e0ac4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/shinkai-bin/shinkai-node/src/managers/model_capabilities_manager.rs b/shinkai-bin/shinkai-node/src/managers/model_capabilities_manager.rs
@@ -289,6 +289,7 @@ impl ModelCapabilitiesManager {
                 if ollama.model_type.starts_with("mistral")
                     || ollama.model_type.starts_with("llama2")
                     || ollama.model_type.starts_with("llama3")
+                    || ollama.model_type.starts_with("gemma2")
                     || ollama.model_type.starts_with("wizardlm2")
                     || ollama.model_type.starts_with("starling-lm")
                     || ollama.model_type.starts_with("neural-chat")
@@ -399,6 +400,7 @@ impl ModelCapabilitiesManager {
                     model_type if model_type.starts_with("dolphin-llama3") => 8_000,
                     model_type if model_type.starts_with("command-r-plus") => 128_000,
                     model_type if model_type.starts_with("codestral") => 32_000,
+                    model_type if model_type.starts_with("gemma2") => 8_000,
                     model_type if model_type.starts_with("qwen2:0.5b") => 32_000,
                     model_type if model_type.starts_with("qwen2:1.5b") => 32_000,
                     model_type if model_type.starts_with("qwen2:7b") => 128_000,