diff --git a/modules/RWKV.py b/modules/RWKV.py new file mode 100644 index 00000000..88f1ec23 --- /dev/null +++ b/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) diff --git a/modules/models.py b/modules/models.py index 1264a58c..955ade0b 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,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" diff --git a/modules/shared.py b/modules/shared.py index f0408068..ec1bd521 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..1324c8b8 100644 --- a/modules/text_generation.py +++ b/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") 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