소스 검색

Add RWKV pip package

oobabooga 2 년 전
부모
커밋
9c86a1cd4a
2개의 변경된 파일2개의 추가작업 그리고 2개의 파일을 삭제
  1. 1 2
      modules/RWKV.py
  2. 1 0
      requirements.txt

+ 1 - 2
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

+ 1 - 0
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