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

19
modules/RWKV.py
View file

@ -16,11 +16,20 @@ os.environ["RWKV_CUDA_ON"] = '0' # '1' : use CUDA kernel for seq mode (much fas
from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS
class RWKVModel:
def __init__(self):
pass
def load_RWKV_model(path):
print(f'strategy={"cpu" if shared.args.cpu else "cuda"} {"fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16"}')
@classmethod
def from_pretrained(self, path, dtype="fp16", device="cuda"):
tokenizer_path = Path(f"{path.parent}/20B_tokenizer.json")
model = RWKV(model=path.as_posix(), strategy=f'{"cpu" if shared.args.cpu else "cuda"} {"fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16"}')
pipeline = PIPELINE(model, Path("models/20B_tokenizer.json").as_posix())
model = RWKV(model=path.as_posix(), strategy=f'{device} {dtype}')
pipeline = PIPELINE(model, tokenizer_path.as_posix())
return pipeline
result = self()
result.model = pipeline
return result
def generate(self, context, **kwargs):
return self.model.generate(context, **kwargs)