AI/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-10-01 01:26:03 -04:00
Code Issues Packages Projects Releases Wiki Activity

Don't use low_cpu_mem_usage and device_map together

Browse Source
This commit is contained in:
oobabooga 2023-01-31 13:24:05 -03:00
parent 824329749d
commit bb77f20a6c
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
server.py
View File

@ -87,10 +87,11 @@ def load_model(model_name):
model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
# Custom
else:
settings = ["low_cpu_mem_usage=True"]
command = "AutoModelForCausalLM.from_pretrained"
settings = []
if args.cpu:
settings.append("low_cpu_mem_usage=True")
settings.append("torch_dtype=torch.float32")
else:
settings.append("device_map='auto'")
@ -374,7 +375,7 @@ if args.chat or args.cai_chat:
reply = reply[idx + 1 + len(apply_extensions(f"{current}:", "bot_prefix")):]
else:
reply = reply[idx + 1 + len(f"{current}:"):]
if check:
reply = reply.split('\n')[0].strip()
else: