From 680a52f2f2cd98547bd78424289acaea8b77ab2b Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Mon, 26 Feb 2024 11:20:39 +0000 Subject: [PATCH] Fix GPT2 detection. --- server/text_generation_server/models/__init__.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py index 8edf0677..22cd0f57 100644 --- a/server/text_generation_server/models/__init__.py +++ b/server/text_generation_server/models/__init__.py @@ -181,7 +181,11 @@ def get_model( trust_remote_code=trust_remote_code, ) - if model_type in {"gpt_bigcode", "gpt2"}: + if ( + model_type == "gpt_bigcode" + or model_type == "gpt2" + and model_id.startswith("bigcode/") + ): if FLASH_ATTENTION: return FlashSantacoderSharded( model_id,