zhengjl/text-generation-webui
Public Access
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Make the code more like PEP8

Browse Source
This commit is contained in:
oobabooga
2023-04-07 00:00:16 -03:00
parent 8e4e8dd741
commit 7b1b1d8e39
30 changed files with 314 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
modules/LoRA.py
View File

@@ -13,6 +13,7 @@ def reload_model():
clear_torch_cache()
shared.model, shared.tokenizer = load_model(shared.model_name)
def add_lora_to_model(lora_name):
# If a LoRA had been previously loaded, or if we want
@@ -27,10 +28,10 @@ def add_lora_to_model(lora_name):
if not shared.args.cpu:
params['dtype'] = shared.model.dtype
if hasattr(shared.model, "hf_device_map"):
params['device_map'] = {"base_model.model."+k: v for k, v in shared.model.hf_device_map.items()}
params['device_map'] = {"base_model.model." + k: v for k, v in shared.model.hf_device_map.items()}
elif shared.args.load_in_8bit:
params['device_map'] = {'': 0}
shared.model = PeftModel.from_pretrained(shared.model, Path(f"{shared.args.lora_dir}/{lora_name}"), **params)
if not shared.args.load_in_8bit and not shared.args.cpu:
shared.model.half()