support granite-code models

ayb · Jun 4, 2024 · 6a9e416 · 6a9e416
1 parent 672f4cf
commit 6a9e416
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/app/lib/llm_clients/client.rb b/app/lib/llm_clients/client.rb
@@ -21,6 +21,7 @@ class Client
     MODEL_TEMPLATE_MAP = {
       "^gemma:" => { prefix: "<start_of_turn>user\n", suffix: "<end_of_turn>\n<start_of_turn>model" },
       "^mi[xs]tral:" => { prefix: "<s>[INST]", suffix: "[/INST] " },
+      "^granite-code:" => { prefix: "<|user|>", suffix: "<|assistant|>" },
     }.freeze
     TIMEOUT_RETRIES = 3