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

Small fixes in the weights loading logic.

Browse Source
This commit is contained in:
Nicolas Patry 2024-02-26 17:32:42 +00:00
parent 915e5f088c
commit e69e68c8ea
2 changed files with 4 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
server/text_generation_server/cli.py
View File

@ -194,16 +194,12 @@ def download_weights(
if not extension == ".safetensors" or not auto_convert:
raise e
elif (Path(model_id) / "medusa_lm_head.pt").exists():
elif (Path(model_id) / "medusa_lm_head.safetensors").exists():
# Try to load as a local Medusa model
try:
import json
medusa_head = Path(model_id) / "medusa_lm_head.pt"
if auto_convert:
medusa_sf = Path(model_id) / "medusa_lm_head.safetensors"
if not medusa_sf.exists():
utils.convert_files([Path(medusa_head)], [medusa_sf], [])
medusa_head = Path(model_id) / "medusa_lm_head.safetensors"
medusa_config = Path(model_id) / "config.json"
with open(medusa_config, "r") as f:
config = json.load(f)

2
server/text_generation_server/utils/hub.py
View File

@ -40,6 +40,7 @@ def _weight_hub_files_from_model_info(
and "arguments" not in s.rfilename
and "args" not in s.rfilename
and "training" not in s.rfilename
and "medusa_lm_head" not in s.rfilename
]
@ -56,6 +57,7 @@ def _weight_files_from_dir(d: Path, extension: str) -> List[str]:
and "args" not in f
and "adapter" not in f
and "training" not in f
and "medusa_lm_head" not in f
]
return filenames