huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-27 13:02:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: fix local loading for .bin models (#1419)

Browse Source
This commit is contained in:
OlivierDehaene 2024-01-09 15:21:00 +01:00 committed by GitHub
parent 3f9b3f4539
commit 564f2a3b75
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/text_generation_server/cli.py
View File

@ -198,7 +198,7 @@ def download_weights(
if not extension == ".safetensors" or not auto_convert: if not extension == ".safetensors" or not auto_convert:
raise e raise e
else: elif (Path(model_id) / "adapter_config.json").exists():
# Try to load as a local PEFT model # Try to load as a local PEFT model
try: try:
utils.download_and_unload_peft( utils.download_and_unload_peft(

5
server/text_generation_server/utils/peft.py
View File

@ -10,8 +10,7 @@ from peft import AutoPeftModelForCausalLM, AutoPeftModelForSeq2SeqLM
def download_and_unload_peft(model_id, revision, trust_remote_code): def download_and_unload_peft(model_id, revision, trust_remote_code):
torch_dtype = torch.float16 torch_dtype = torch.float16
logger.info("Peft model detected.") logger.info("Trying to load a Peft model. It might take a while without feedback")
logger.info("Loading the model it might take a while without feedback")
try: try:
model = AutoPeftModelForCausalLM.from_pretrained( model = AutoPeftModelForCausalLM.from_pretrained(
model_id, model_id,
@ -28,7 +27,7 @@ def download_and_unload_peft(model_id, revision, trust_remote_code):
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
logger.info(f"Loaded.") logger.info("Peft model detected.")
logger.info(f"Merging the lora weights.") logger.info(f"Merging the lora weights.")
base_model_id = model.peft_config["default"].base_model_name_or_path base_model_id = model.peft_config["default"].base_model_name_or_path