Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-22 08:07:56 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add RWKV pip package

Browse Source
This commit is contained in:
oobabooga 2023-03-01 11:42:49 -03:00
parent 6837d4d72a
commit 9c86a1cd4a
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
modules/RWKV.py
View File

@ -13,7 +13,6 @@ np.set_printoptions(precision=4, suppress=True, linewidth=200)
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0' # '1' : use CUDA kernel for seq mode (much faster)
import repositories.ChatRWKV.v2.rwkv as rwkv
from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS
@ -22,6 +21,6 @@ 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"}')
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("repositories/ChatRWKV/20B_tokenizer.json").as_posix())
pipeline = PIPELINE(model, Path("models/20B_tokenizer.json").as_posix())
return pipeline

1
requirements.txt
View File

@ -3,5 +3,6 @@ bitsandbytes==0.37.0
flexgen==0.1.6
gradio==3.18.0
numpy
rwkv==0.0.5
safetensors==0.2.8
git+https://github.com/huggingface/transformers