dwolfe66/text-generation-webui-space

1import os2from pathlib import Path3 4import numpy as np5from tokenizers import Tokenizer6 7import modules.shared as shared8from modules.callbacks import Iteratorize9 10np.set_printoptions(precision=4, suppress=True, linewidth=200)11 12os.environ['RWKV_JIT_ON'] = '1'13os.environ["RWKV_CUDA_ON"] = '1' if shared.args.rwkv_cuda_on else '0' # use CUDA kernel for seq mode (much faster)14 15from rwkv.model import RWKV16from rwkv.utils import PIPELINE, PIPELINE_ARGS17 18 19class RWKVModel:20    def __init__(self):21        pass22 23    @classmethod24    def from_pretrained(self, path, dtype="fp16", device="cuda"):25        tokenizer_path = Path(f"{path.parent}/20B_tokenizer.json")26 27        if shared.args.rwkv_strategy is None:28            model = RWKV(model=str(path), strategy=f'{device} {dtype}')29        else:30            model = RWKV(model=str(path), strategy=shared.args.rwkv_strategy)31        pipeline = PIPELINE(model, str(tokenizer_path))32 33        result = self()34        result.pipeline = pipeline35        return result36 37    def generate(self, context="", token_count=20, temperature=1, top_p=1, top_k=50, alpha_frequency=0.1, alpha_presence=0.1, token_ban=[0], token_stop=[], callback=None):38        args = PIPELINE_ARGS(39            temperature = temperature,40            top_p = top_p,41            top_k = top_k,42            alpha_frequency = alpha_frequency, # Frequency Penalty (as in GPT-3)43            alpha_presence = alpha_presence, # Presence Penalty (as in GPT-3)44            token_ban = token_ban, # ban the generation of some tokens45            token_stop = token_stop46        )47 48        return context+self.pipeline.generate(context, token_count=token_count, args=args, callback=callback)49 50    def generate_with_streaming(self, **kwargs):51        with Iteratorize(self.generate, kwargs, callback=None) as generator:52            reply = kwargs['context']53            for token in generator:54                reply += token55                yield reply56 57class RWKVTokenizer:58    def __init__(self):59        pass60 61    @classmethod62    def from_pretrained(self, path):63        tokenizer_path = path / "20B_tokenizer.json"64        tokenizer = Tokenizer.from_file(str(tokenizer_path))65 66        result = self()67        result.tokenizer = tokenizer68        return result69 70    def encode(self, prompt):71        return self.tokenizer.encode(prompt).ids72 73    def decode(self, ids):74        return self.tokenizer.decode(ids)75