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

Remove RWKV loader (#5130)

Browse source
This commit is contained in:
oobabooga 2023-12-31 02:01:40 -03:00 • committed by GitHub
parent 0e54a09bcb
commit 2734ce3e4c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 3 additions and 189 deletions
Download patch file Download diff file Expand all files Collapse all files

18
modules/models.py
View file

@ -65,7 +65,6 @@ def load_model(model_name, loader=None):
'GPTQ-for-LLaMa': GPTQ_loader,
'llama.cpp': llamacpp_loader,
'llamacpp_HF': llamacpp_HF_loader,
'RWKV': RWKV_loader,
'ExLlamav2': ExLlamav2_loader,
'ExLlamav2_HF': ExLlamav2_HF_loader,
'ctransformers': ctransformers_loader,
@ -405,23 +404,6 @@ def HQQ_loader(model_name):
return model
def RWKV_loader(model_name):
'''
This loader is not currently maintained as RWKV can now be loaded
through the transformers library.
'''
from modules.RWKV import RWKVModel, RWKVTokenizer
model = RWKVModel.from_pretrained(
Path(f'{shared.args.model_dir}/{model_name}'),
dtype="fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16",
device="cpu" if shared.args.cpu else "xpu" if is_xpu_available() else "cuda"
)
tokenizer = RWKVTokenizer.from_pretrained(Path(shared.args.model_dir))
return model, tokenizer
def get_max_memory_dict():
max_memory = {}
max_cpu_memory = shared.args.cpu_memory.strip() if shared.args.cpu_memory is not None else '99GiB'