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

Adapt to the new model names

Browse source
This commit is contained in:
oobabooga 2023-03-29 21:47:36 -03:00
parent 0345e04249
commit 1cb9246160
6 changed files with 18 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

4
modules/models.py
View file

@ -41,7 +41,7 @@ def load_model(model_name):
print(f"Loading {model_name}...")
t0 = time.time()
shared.is_RWKV = model_name.lower().startswith('rwkv-')
shared.is_RWKV = 'rwkv-' in model_name.lower()
# Default settings
if not any([shared.args.cpu, shared.args.load_in_8bit, shared.args.wbits, shared.args.auto_devices, shared.args.disk, shared.args.gpu_memory is not None, shared.args.cpu_memory is not None, shared.args.deepspeed, shared.args.flexgen, shared.is_RWKV]):
@ -159,7 +159,7 @@ def load_model(model_name):
model = AutoModelForCausalLM.from_pretrained(checkpoint, **params)
# Loading the tokenizer
if shared.model_name.lower().startswith(('gpt4chan', 'gpt-4chan', '4chan')) and Path(f"{shared.args.model_dir}/gpt-j-6B/").exists():
if any((k in shared.model_name.lower() for k in ['gpt4chan', 'gpt-4chan'])) and Path(f"{shared.args.model_dir}/gpt-j-6B/").exists():
tokenizer = AutoTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/gpt-j-6B/"))
else:
tokenizer = AutoTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/{shared.model_name}/"))