From c0e9b46795153c452e05dc0b3b7cf23a3cd7b188 Mon Sep 17 00:00:00 2001 From: bmullick-amd Date: Mon, 2 Dec 2024 13:23:46 -0800 Subject: [PATCH] Update registry.py --- vllm/model_executor/models/registry.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index f17f31d20857..6bdd9a823f82 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -90,6 +90,9 @@ "BartModel": ("bart", "BartForConditionalGeneration"), "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"), "Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"), # noqa: E501 + "T5Model": ("t5", "T5ForConditionalGeneration"), + "T5ForConditionalGeneration": ("t5", "T5ForConditionalGeneration"), + "T5WithLMHeadModel": ("t5", "T5ForConditionalGeneration") } _EMBEDDING_MODELS = { @@ -484,4 +487,4 @@ def _run() -> None: if __name__ == "__main__": - _run() \ No newline at end of file + _run()