From 6a25c14f465a835fc006ecf48debed511237c90d Mon Sep 17 00:00:00 2001 From: Traun Leyden Date: Tue, 14 Nov 2023 15:22:03 +0100 Subject: [PATCH] Load local PEFT model --- server/text_generation_server/cli.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/server/text_generation_server/cli.py b/server/text_generation_server/cli.py index 301acb6b..c8f40c46 100644 --- a/server/text_generation_server/cli.py +++ b/server/text_generation_server/cli.py @@ -150,6 +150,16 @@ def download_weights( if not extension == ".safetensors" or not auto_convert: raise e + # Try to load as a local PEFT model + try: + utils.download_and_unload_peft( + model_id, revision, trust_remote_code=trust_remote_code + ) + utils.weight_files(model_id, revision, extension) + return + except (utils.LocalEntryNotFoundError, utils.EntryNotFoundError): + pass + # Try to see if there are local pytorch weights try: # Get weights for a local model, a hub cached model and inside the WEIGHTS_CACHE_OVERRIDE