diff --git a/modules/RWKV.py b/modules/RWKV.py index 3d6a2be5..b7388ea7 100644 --- a/modules/RWKV.py +++ b/modules/RWKV.py @@ -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 diff --git a/requirements.txt b/requirements.txt index b333ffba..7dcd720a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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