From ebc64a408c1c36df98a4487eefc10a3876500e6c Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Mon, 27 Feb 2023 23:03:35 -0300
Subject: [PATCH] RWKV support prototype

---
 modules/models.py          | 28 +++++++++++++++++++++++++++-
 modules/shared.py          |  1 +
 modules/text_generation.py | 14 ++++++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/modules/models.py b/modules/models.py
index 1264a58c..9ce94f6b 100644
--- a/modules/models.py
+++ b/modules/models.py
@@ -38,8 +38,10 @@ def load_model(model_name):
     print(f"Loading {model_name}...")
     t0 = time.time()
 
+    shared.is_RWKV = model_name.lower().startswith('rwkv-')
+
     # Default settings
-    if not (shared.args.cpu or shared.args.load_in_8bit or shared.args.auto_devices or shared.args.disk or shared.args.gpu_memory is not None or shared.args.cpu_memory is not None or shared.args.deepspeed or shared.args.flexgen):
+    if not (shared.args.cpu or shared.args.load_in_8bit or shared.args.auto_devices or shared.args.disk or shared.args.gpu_memory is not None or shared.args.cpu_memory is not None or shared.args.deepspeed or shared.args.flexgen or shared.is_RWKV):
         if any(size in shared.model_name.lower() for size in ('13b', '20b', '30b')):
             model = AutoModelForCausalLM.from_pretrained(Path(f"models/{shared.model_name}"), device_map='auto', load_in_8bit=True)
         else:
@@ -75,6 +77,30 @@ def load_model(model_name):
         model.module.eval() # Inference
         print(f"DeepSpeed ZeRO-3 is enabled: {is_deepspeed_zero3_enabled()}")
 
+    # RMKV model (not on HuggingFace)
+    elif shared.is_RWKV:
+        import types
+        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)
+
+        from rwkv.model import RWKV
+        from rwkv.utils import PIPELINE, PIPELINE_ARGS
+
+        model = RWKV(model='models/RWKV-4-Pile-169M-20220807-8023.pth', strategy='cuda fp16')
+
+        out, state = model.forward([187, 510, 1563, 310, 247], None)   # use 20B_tokenizer.json
+        print(out.detach().cpu().numpy())                   # get logits
+        out, state = model.forward([187, 510], None)
+        out, state = model.forward([1563], state)           # RNN has state (use deepcopy if you want to clone it)
+        out, state = model.forward([310, 247], state)
+        print(out.detach().cpu().numpy())                   # same result as above
+
+        pipeline = PIPELINE(model, "20B_tokenizer.json")
+
+        return pipeline, None
+
     # Custom
     else:
         command = "AutoModelForCausalLM.from_pretrained"
diff --git a/modules/shared.py b/modules/shared.py
index d59cee99..b28f8c5f 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -5,6 +5,7 @@ tokenizer = None
 model_name = ""
 soft_prompt_tensor = None
 soft_prompt = False
+is_RWKV = False
 
 # Chat variables
 history = {'internal': [], 'visible': []}
diff --git a/modules/text_generation.py b/modules/text_generation.py
index 9c8674d2..ebe6ed35 100644
--- a/modules/text_generation.py
+++ b/modules/text_generation.py
@@ -6,6 +6,7 @@ import numpy as np
 import torch
 import transformers
 from tqdm import tqdm
+from rwkv.utils import PIPELINE, PIPELINE_ARGS
 
 import modules.shared as shared
 from modules.extensions import apply_extensions
@@ -80,6 +81,19 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
     if not shared.args.cpu:
         torch.cuda.empty_cache()
 
+    if shared.is_RWKV:
+        def my_print(s):
+            print(s, end='', flush=True)
+        args = PIPELINE_ARGS(temperature = temperature, top_p = top_p,
+                             alpha_frequency = 0.25, # Frequency Penalty (as in GPT-3)
+                             alpha_presence = 0.25, # Presence Penalty (as in GPT-3)
+                             token_ban = [0], # ban the generation of some tokens
+                             token_stop = []) # stop generation whenever you see any token here
+        reply = question + shared.model.generate(question, token_count=max_new_tokens, args=args, callback=None)
+        print(formatted_outputs(reply, None))
+        yield formatted_outputs(reply, None)
+    return formatted_outputs(reply, None)
+
     original_question = question
     if not (shared.args.chat or shared.args.cai_chat):
         question = apply_extensions(question, "input")