huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix GPT2 detection.

Browse Source
This commit is contained in:
Nicolas Patry 2024-02-26 11:20:39 +00:00
parent ed95f1982d
commit 680a52f2f2
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
server/text_generation_server/models/__init__.py
View File

@ -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,