[Model] Support Cohere2ForCausalLM (Cohere R7B) (#11203)

2024-12-16 11:56:19 +02:00
parent b3b1526f03
commit bddbbcb132
5 changed files with 26 additions and 4 deletions
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -41,6 +41,7 @@ _TEXT_GENERATION_MODELS = {
    "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
    # ChatGLMModel supports multimodal
    "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
+    "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
    "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
    "DeciLMForCausalLM": ("decilm", "DeciLMForCausalLM"),
    "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),