diff --git a/docs/source/models/supported_models.md b/docs/source/models/supported_models.md index 2fb969ea85f1..368c903c3252 100644 --- a/docs/source/models/supported_models.md +++ b/docs/source/models/supported_models.md @@ -233,9 +233,9 @@ See [this page](#generative-models) for more information on how to use generativ * `facebook/bart-base`, `facebook/bart-large-cnn`, etc. * * -- * `ChatGLMModel` +- * `ChatGLMModel`, `ChatGLMForConditionalGeneration` * ChatGLM - * `THUDM/chatglm2-6b`, `THUDM/chatglm3-6b`, etc. + * `THUDM/chatglm2-6b`, `THUDM/chatglm3-6b`, `ShieldLM-6B-chatglm3`, etc. * ✅︎ * ✅︎ - * `CohereForCausalLM`, `Cohere2ForCausalLM` diff --git a/tests/models/registry.py b/tests/models/registry.py index e61cbc5756f6..10b93460c56b 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -124,6 +124,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "BloomForCausalLM": _HfExamplesInfo("bigscience/bloomz-1b1"), "ChatGLMModel": _HfExamplesInfo("THUDM/chatglm3-6b", trust_remote_code=True), + "ChatGLMForConditionalGeneration": _HfExamplesInfo("thu-coai/ShieldLM-6B-chatglm3", # noqa: E501 + trust_remote_code=True), "CohereForCausalLM": _HfExamplesInfo("CohereForAI/c4ai-command-r-v01", trust_remote_code=True), "Cohere2ForCausalLM": _HfExamplesInfo("CohereForAI/c4ai-command-r7b-12-2024", # noqa: E501 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 3abbb1f0c3b6..43ff892349e2 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -43,6 +43,7 @@ _TEXT_GENERATION_MODELS = { "BambaForCausalLM": ("bamba", "BambaForCausalLM"), "BloomForCausalLM": ("bloom", "BloomForCausalLM"), "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"), + "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"), "CohereForCausalLM": ("commandr", "CohereForCausalLM"), "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"), "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),