From 9c86a1cd4ae3ab9a4f1e3f8ee093a5f6964815c2 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 1 Mar 2023 11:42:49 -0300
Subject: [PATCH] Add RWKV pip package

---
 modules/RWKV.py  | 3 +--
 requirements.txt | 1 +
 2 files changed, 2 insertions(+), 2 deletions(-)

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