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

Reorganize things

Browse Source
This commit is contained in:
oobabooga 2023-01-15 23:01:51 -03:00
parent a3bb18002a
commit b55486fa00
1 changed files with 3 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
server.py
View File

@ -74,28 +74,21 @@ def load_model(model_name):
# Custom
else:
settings = ["low_cpu_mem_usage=True"]
cuda = ""
command = "AutoModelForCausalLM.from_pretrained"
if args.cpu:
settings.append("torch_dtype=torch.float32")
else:
settings.append("device_map='auto'")
if args.max_gpu_memory is not None:
settings.append(f"max_memory={{0: '{args.max_gpu_memory}GiB', 'cpu': '99GiB'}}")
settings.append("device_map='auto'")
settings.append("torch_dtype=torch.float16")
elif args.load_in_8bit:
settings.append("device_map='auto'")
if args.load_in_8bit:
settings.append("load_in_8bit=True")
else:
settings.append("torch_dtype=torch.float16")
if args.auto_devices:
settings.append("device_map='auto'")
else:
cuda = ".cuda()"
settings = ', '.join(list(set(settings)))
command = f"{command}(Path(f'models/{model_name}'), {settings}){cuda}"
command = f"{command}(Path(f'models/{model_name}'), {settings})"
model = eval(command)
# Loading the tokenizer