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

Add RWKVModel class

Browse source
This commit is contained in:
oobabooga 2023-03-01 12:08:55 -03:00
parent 9c86a1cd4a
commit 659bb76722
2 changed files with 18 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

6
modules/models.py
View file

@ -79,9 +79,11 @@ def load_model(model_name):
# RMKV model (not on HuggingFace)
elif shared.is_RWKV:
from modules.RWKV import load_RWKV_model
from modules.RWKV import RWKVModel
return load_RWKV_model(Path(f'models/{model_name}')), None
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")
return model, None
# Custom
else: