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

Add RWKV tokenizer

Browse source
This commit is contained in:
oobabooga 2023-03-06 08:45:49 -03:00
parent c855b828fe
commit e91f4bc25a
3 changed files with 34 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

5
modules/models.py
View file

@ -79,11 +79,12 @@ def load_model(model_name):
# RMKV model (not on HuggingFace)
elif shared.is_RWKV:
from modules.RWKV import RWKVModel
from modules.RWKV import RWKVModel, RWKVTokenizer
model = RWKVModel.from_pretrained(Path(f'models/{model_name}'), dtype="fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16", device="cpu" if shared.args.cpu else "cuda")
tokenizer = RWKVTokenizer.from_pretrained(Path('models'))
return model, None
return model, tokenizer
# Custom
else: