forestcalled/text-generation-webui

1'''2This loader is not currently maintained as RWKV can now be loaded3through the transformers library.4'''5 6import copy7import os8from pathlib import Path9 10import numpy as np11from tokenizers import Tokenizer12from transformers import is_torch_xpu_available13 14import modules.shared as shared15from modules.callbacks import Iteratorize16 17np.set_printoptions(precision=4, suppress=True, linewidth=200)18 19os.environ['RWKV_JIT_ON'] = '1'20os.environ["RWKV_CUDA_ON"] = '1' if shared.args.rwkv_cuda_on else '0'  # use CUDA kernel for seq mode (much faster)21 22from rwkv.model import RWKV23from rwkv.utils import PIPELINE, PIPELINE_ARGS24 25 26class RWKVModel:27    def __init__(self):28        pass29 30    @classmethod31    def from_pretrained(self, path, dtype="bf16" if is_torch_xpu_available() else "fp16", device="xpu" if is_torch_xpu_available() else "cuda"):32        tokenizer_path = Path(f"{path.parent}/20B_tokenizer.json")33        if shared.args.rwkv_strategy is None:34            model = RWKV(model=str(path), strategy=f'{device} {dtype}')35        else:36            model = RWKV(model=str(path), strategy=shared.args.rwkv_strategy)37 38        pipeline = PIPELINE(model, str(tokenizer_path))39        result = self()40        result.pipeline = pipeline41        result.model = model42        result.cached_context = ""43        result.cached_model_state = None44        result.cached_output_logits = None45        return result46 47    def generate(self, prompt, state, callback=None):48        args = PIPELINE_ARGS(49            temperature=state['temperature'],50            top_p=state['top_p'],51            top_k=state['top_k'],52            alpha_frequency=0.1,  # Frequency Penalty (as in GPT-3)53            alpha_presence=0.1,  # Presence Penalty (as in GPT-3)54            token_ban=[0],  # ban the generation of some tokens55            token_stop=[]56        )57 58        if self.cached_context != "":59            if prompt.startswith(self.cached_context):60                prompt = prompt[len(self.cached_context):]61            else:62                self.cached_context = ""63                self.cached_model_state = None64                self.cached_output_logits = None65 66        # out = self.pipeline.generate(prompt, token_count=state['max_new_tokens'], args=args, callback=callback)67        out = self.generate_from_cached_state(prompt, token_count=state['max_new_tokens'], args=args, callback=callback)68        return out69 70    def generate_with_streaming(self, *args, **kwargs):71        with Iteratorize(self.generate, args, kwargs, callback=None) as generator:72            reply = ''73            for token in generator:74                reply += token75                yield reply76 77    # Similar to the PIPELINE.generate, but lets us maintain the cached_model_state78    def generate_from_cached_state(self, ctx="", token_count=20, args=None, callback=None):79        all_tokens = []80        out_str = ''81        occurrence = {}82        state = copy.deepcopy(self.cached_model_state) if self.cached_model_state is not None else None83 84        # if we ended up with an empty context, just reuse the cached logits85        # this can happen if a user undoes a message and then sends the exact message again86        # in that case the full context ends up being the same as the cached_context, so the remaining context is empty.87        if ctx == "":88            out = self.cached_output_logits89 90        token = None91        for i in range(token_count):92            # forward93            tokens = self.pipeline.encode(ctx) if i == 0 else [token]94            while len(tokens) > 0:95                out, state = self.model.forward(tokens[:args.chunk_len], state)96                tokens = tokens[args.chunk_len:]97            if i == 0:98                begin_token = len(all_tokens)99                last_token_posi = begin_token100            # cache the model state after scanning the context101            # we don't cache the state after processing our own generated tokens because102            # the output string might be post-processed arbitrarily. Therefore, what's fed into the model103            # on the next round of chat might be slightly different what what it output on the previous round104            if i == 0:105                self.cached_context += ctx106                self.cached_model_state = copy.deepcopy(state)107                self.cached_output_logits = copy.deepcopy(out)108 109            # adjust probabilities110            for n in args.token_ban:111                out[n] = -float('inf')112 113            for n in occurrence:114                out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)115 116            # sampler117            token = self.pipeline.sample_logits(out, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k)118            if token in args.token_stop:119                break120 121            all_tokens += [token]122            if token not in occurrence:123                occurrence[token] = 1124            else:125                occurrence[token] += 1126 127            # output128            tmp = self.pipeline.decode(all_tokens[last_token_posi:])129            if '\ufffd' not in tmp:  # is valid utf-8 string?130                if callback:131                    callback(tmp)132 133                out_str += tmp134                last_token_posi = begin_token + i + 1135        return out_str136 137 138class RWKVTokenizer:139    def __init__(self):140        pass141 142    @classmethod143    def from_pretrained(self, path):144        tokenizer_path = path / "20B_tokenizer.json"145        tokenizer = Tokenizer.from_file(str(tokenizer_path))146        result = self()147        result.tokenizer = tokenizer148        return result149 150    def encode(self, prompt):151        return self.tokenizer.encode(prompt).ids152 153    def decode(self, ids):154        return self.tokenizer.decode(ids)155