forestcalled/text-generation-webui

1import re2from functools import partial3 4import numpy as np5import torch6 7from modules import RoPE, shared8from modules.callbacks import Iteratorize9from modules.logging_colors import logger10from modules.text_generation import get_max_prompt_length11 12try:13    import llama_cpp14except:15    llama_cpp = None16 17try:18    import llama_cpp_cuda19except:20    llama_cpp_cuda = None21 22try:23    import llama_cpp_cuda_tensorcores24except:25    llama_cpp_cuda_tensorcores = None26 27 28def llama_cpp_lib():29    if shared.args.cpu and llama_cpp is not None:30        return llama_cpp31    elif shared.args.tensorcores and llama_cpp_cuda_tensorcores is not None:32        return llama_cpp_cuda_tensorcores33    elif llama_cpp_cuda is not None:34        return llama_cpp_cuda35    else:36        return llama_cpp37 38 39def ban_eos_logits_processor(eos_token, input_ids, logits):40    logits[eos_token] = -float('inf')41    return logits42 43 44def custom_token_ban_logits_processor(token_ids, input_ids, logits):45    for token_id in token_ids:46        logits[token_id] = -float('inf')47 48    return logits49 50 51class LlamaCppModel:52    def __init__(self):53        self.initialized = False54        self.grammar_string = ''55        self.grammar = None56 57    def __del__(self):58        del self.model59 60    @classmethod61    def from_pretrained(self, path):62 63        Llama = llama_cpp_lib().Llama64        LlamaCache = llama_cpp_lib().LlamaCache65 66        result = self()67        cache_capacity = 068        if shared.args.cache_capacity is not None:69            if 'GiB' in shared.args.cache_capacity:70                cache_capacity = int(re.sub('[a-zA-Z]', '', shared.args.cache_capacity)) * 1000 * 1000 * 100071            elif 'MiB' in shared.args.cache_capacity:72                cache_capacity = int(re.sub('[a-zA-Z]', '', shared.args.cache_capacity)) * 1000 * 100073            else:74                cache_capacity = int(shared.args.cache_capacity)75 76        if cache_capacity > 0:77            logger.info("Cache capacity is " + str(cache_capacity) + " bytes")78 79        if shared.args.tensor_split is None or shared.args.tensor_split.strip() == '':80            tensor_split_list = None81        else:82            tensor_split_list = [float(x) for x in shared.args.tensor_split.strip().split(",")]83 84        params = {85            'model_path': str(path),86            'n_ctx': shared.args.n_ctx,87            'n_threads': shared.args.threads or None,88            'n_threads_batch': shared.args.threads_batch or None,89            'n_batch': shared.args.n_batch,90            'use_mmap': not shared.args.no_mmap,91            'use_mlock': shared.args.mlock,92            'mul_mat_q': not shared.args.no_mul_mat_q,93            'numa': shared.args.numa,94            'n_gpu_layers': shared.args.n_gpu_layers,95            'rope_freq_base': RoPE.get_rope_freq_base(shared.args.alpha_value, shared.args.rope_freq_base),96            'tensor_split': tensor_split_list,97            'rope_freq_scale': 1.0 / shared.args.compress_pos_emb,98            'offload_kqv': not shared.args.no_offload_kqv99        }100 101        result.model = Llama(**params)102        if cache_capacity > 0:103            result.model.set_cache(LlamaCache(capacity_bytes=cache_capacity))104 105        # This is ugly, but the model and the tokenizer are the same object in this library.106        return result, result107 108    def encode(self, string):109        if type(string) is str:110            string = string.encode()111 112        return self.model.tokenize(string)113 114    def decode(self, ids, **kwargs):115        return self.model.detokenize(ids).decode('utf-8')116 117    def get_logits(self, tokens):118        self.model.reset()119        self.model.eval(tokens)120        logits = self.model._scores121        logits = np.expand_dims(logits, 0)  # batch dim is expected122        return torch.tensor(logits, dtype=torch.float32)123 124    def load_grammar(self, string):125        if string != self.grammar_string:126            self.grammar_string = string127            if string.strip() != '':128                self.grammar = llama_cpp_lib().LlamaGrammar.from_string(string)129            else:130                self.grammar = None131 132    def generate(self, prompt, state, callback=None):133        LogitsProcessorList = llama_cpp_lib().LogitsProcessorList134        prompt = prompt if type(prompt) is str else prompt.decode()135 136        # Handle truncation137        prompt = self.encode(prompt)138        prompt = prompt[-get_max_prompt_length(state):]139        prompt = self.decode(prompt)140 141        self.load_grammar(state['grammar_string'])142        logit_processors = LogitsProcessorList()143        if state['ban_eos_token']:144            logit_processors.append(partial(ban_eos_logits_processor, self.model.token_eos()))145 146        if state['custom_token_bans']:147            to_ban = [int(x) for x in state['custom_token_bans'].split(',')]148            if len(to_ban) > 0:149                logit_processors.append(partial(custom_token_ban_logits_processor, to_ban))150 151        completion_chunks = self.model.create_completion(152            prompt=prompt,153            max_tokens=state['max_new_tokens'],154            temperature=state['temperature'],155            top_p=state['top_p'],156            min_p=state['min_p'],157            typical_p=state['typical_p'],158            frequency_penalty=state['frequency_penalty'],159            presence_penalty=state['presence_penalty'],160            repeat_penalty=state['repetition_penalty'],161            top_k=state['top_k'],162            stream=True,163            seed=int(state['seed']) if state['seed'] != -1 else None,164            tfs_z=state['tfs'],165            mirostat_mode=int(state['mirostat_mode']),166            mirostat_tau=state['mirostat_tau'],167            mirostat_eta=state['mirostat_eta'],168            logits_processor=logit_processors,169            grammar=self.grammar170        )171 172        output = ""173        for completion_chunk in completion_chunks:174            if shared.stop_everything:175                break176 177            text = completion_chunk['choices'][0]['text']178            output += text179            if callback:180                callback(text)181 182        return output183 184    def generate_with_streaming(self, *args, **kwargs):185        with Iteratorize(self.generate, args, kwargs, callback=None) as generator:186            reply = ''187            for token in generator:188                reply += token189                yield reply190