Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-26 01:30:20 +01: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() model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
# Custom # Custom
else: else:
settings = ["low_cpu_mem_usage=True"]
command = "AutoModelForCausalLM.from_pretrained" command = "AutoModelForCausalLM.from_pretrained"
settings = []
if args.cpu: if args.cpu:
settings.append("low_cpu_mem_usage=True")
settings.append("torch_dtype=torch.float32") settings.append("torch_dtype=torch.float32")
else: else:
settings.append("device_map='auto'") 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")):] reply = reply[idx + 1 + len(apply_extensions(f"{current}:", "bot_prefix")):]
else: else:
reply = reply[idx + 1 + len(f"{current}:"):] reply = reply[idx + 1 + len(f"{current}:"):]
if check: if check:
reply = reply.split('\n')[0].strip() reply = reply.split('\n')[0].strip()
else: else: