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

Make model loading more transparent

Browse Source
This commit is contained in:
oobabooga 2023-01-06 01:41:52 -03:00
parent c65bad40dc
commit 285032da36
2 changed files with 9 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
README.md
View File

@ -46,15 +46,11 @@ The files that you need to download and put under `models/model-name` (for insta
## Converting to pytorch ## Converting to pytorch
This webui allows you to switch between different models on the fly, so it must be fast to load the models from disk. The script `convert-to-torch.py` allows you to convert models to .pt format, which is about 10x faster to load:
One way to make this process about 10x faster is to convert the models to pytorch format using the script `convert-to-torch.py`. Create a folder called `torch-dumps` and then make the conversion with:
python convert-to-torch.py models/model-name/ python convert-to-torch.py models/model-name/
The output model will be saved to `torch-dumps/model-name.pt`. This is the default way to load all models except for `gpt-neox-20b`, `opt-13b`, `OPT-13B-Erebus`, `gpt-j-6B`, and `flan-t5`. I don't remember why these models are exceptions. The output model will be saved to `torch-dumps/model-name.pt`. When you load a new model from the webui, it will first look for this .pt file; if it is not found, it will load the model as usual from `models/model-name/`.
If I get enough ⭐s on this repository, I will make the process of loading models saner and more customizable.
## Starting the webui ## Starting the webui

12
server.py
View File

@ -1,3 +1,4 @@
import os
import re import re
import time import time
import glob import glob
@ -20,17 +21,18 @@ model_name = 'galactica-6.7b'
settings_name = "Default" settings_name = "Default"
def load_model(model_name): def load_model(model_name):
print(f"Loading {model_name}") print(f"Loading {model_name}...")
t0 = time.time() t0 = time.time()
if model_name in ['gpt-neox-20b', 'opt-13b', 'OPT-13B-Erebus']:
if os.path.exists(f"torch-dumps/{model_name}.pt"):
print("Loading in .pt format...")
model = torch.load(f"torch-dumps/{model_name}.pt").cuda()
elif model_name in ['gpt-neox-20b', 'opt-13b', 'OPT-13B-Erebus']:
model = AutoModelForCausalLM.from_pretrained(f"models/{model_name}", device_map='auto', load_in_8bit=True) model = AutoModelForCausalLM.from_pretrained(f"models/{model_name}", device_map='auto', load_in_8bit=True)
elif model_name in ['gpt-j-6B']: elif model_name in ['gpt-j-6B']:
model = AutoModelForCausalLM.from_pretrained(f"models/{model_name}", low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda() model = AutoModelForCausalLM.from_pretrained(f"models/{model_name}", low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
elif model_name in ['flan-t5']: elif model_name in ['flan-t5']:
model = T5ForConditionalGeneration.from_pretrained(f"models/{model_name}").cuda() model = T5ForConditionalGeneration.from_pretrained(f"models/{model_name}").cuda()
else:
model = torch.load(f"torch-dumps/{model_name}.pt").cuda()
if model_name in ['gpt4chan_model_float16']: if model_name in ['gpt4chan_model_float16']:
tokenizer = AutoTokenizer.from_pretrained("models/gpt-j-6B/") tokenizer = AutoTokenizer.from_pretrained("models/gpt-j-6B/")