mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Don't require GPT-J to be installed to load gpt4chan

Browse source
This commit is contained in:
oobabooga 2023-01-09 11:39:13 -03:00
parent 6cbfe19c23
commit a751d7e693
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
server.py
View file

@ -47,7 +47,7 @@ def load_model(model_name):
model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=dtype) model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=dtype)
# Loading the tokenizer # Loading the tokenizer
if model_name.lower().startswith('gpt4chan'): if model_name.lower().startswith('gpt4chan') and Path(f"models/gpt-j-6B/").exists():
tokenizer = AutoTokenizer.from_pretrained(Path("models/gpt-j-6B/")) tokenizer = AutoTokenizer.from_pretrained(Path("models/gpt-j-6B/"))
elif model_name in ['flan-t5', 't5-large']: elif model_name in ['flan-t5', 't5-large']:
tokenizer = T5Tokenizer.from_pretrained(Path(f"models/{model_name}/")) tokenizer = T5Tokenizer.from_pretrained(Path(f"models/{model_name}/"))