Merge pull request #149 from oobabooga/RWKV

Add RWKV support

modules/RWKV.py

@@ -0,0 +1,45 @@
+import os
+import time
+import types
+from pathlib import Path
+
+import numpy as np
+import torch
+
+import modules.shared as shared
+
+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
+
+
+class RWKVModel:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def from_pretrained(self, path, dtype="fp16", device="cuda"):
+        tokenizer_path = Path(f"{path.parent}/20B_tokenizer.json")
+
+        model = RWKV(model=path.as_posix(), strategy=f'{device} {dtype}')
+        pipeline = PIPELINE(model, tokenizer_path.as_posix())
+
+        result = self()
+        result.pipeline = pipeline
+        return result
+
+    def generate(self, context, token_count=20, temperature=1, top_p=1, alpha_frequency=0.25, alpha_presence=0.25, token_ban=[0], token_stop=[], callback=None):
+        args = PIPELINE_ARGS(
+            temperature = temperature,
+            top_p = top_p,
+            alpha_frequency = alpha_frequency, # Frequency Penalty (as in GPT-3)
+            alpha_presence = alpha_presence, # Presence Penalty (as in GPT-3)
+            token_ban = token_ban, # ban the generation of some tokens
+            token_stop = token_stop
+        )
+
+        return context+self.pipeline.generate(context, token_count=token_count, args=args, callback=callback)

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,14 @@ 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:
+        from modules.RWKV import RWKVModel
+
+        model = RWKVModel.from_pretrained(Path(f'models/{model_name}'), dtype="fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16", device="cpu" if shared.args.cpu else "cuda")
+
+        return model, None
+
     # Custom
     else:
         command = "AutoModelForCausalLM.from_pretrained"

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': []}

modules/text_generation.py

@@ -5,6 +5,7 @@ import time
 import numpy as np
 import torch
 import transformers
+from rwkv.utils import PIPELINE, PIPELINE_ARGS
 from tqdm import tqdm
 
 import modules.shared as shared
@@ -21,6 +22,9 @@ def get_max_prompt_length(tokens):
     return max_length
 
 def encode(prompt, tokens_to_generate=0, add_special_tokens=True):
+    if shared.is_RWKV:
+        return prompt
+
     input_ids = shared.tokenizer.encode(str(prompt), return_tensors='pt', truncation=True, max_length=get_max_prompt_length(tokens_to_generate), add_special_tokens=add_special_tokens)
     if shared.args.cpu:
         return input_ids
@@ -80,6 +84,17 @@ 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:
+        if shared.args.no_stream:
+            reply = shared.model.generate(question, token_count=max_new_tokens, temperature=temperature, top_p=top_p)
+            yield formatted_outputs(reply, None)
+        else:
+            for i in range(max_new_tokens//8):
+                reply = shared.model.generate(question, token_count=8, temperature=temperature, top_p=top_p)
+                yield formatted_outputs(reply, None)
+                question = reply
+        return formatted_outputs(reply, None)
+
     original_question = question
     if not (shared.args.chat or shared.args.cai_chat):
         question = apply_extensions(question, "input")

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