forestcalled/text-generation-webui

1import traceback2from pathlib import Path3 4import torch5from exllamav2 import (6    ExLlamaV2,7    ExLlamaV2Cache,8    ExLlamaV2Cache_8bit,9    ExLlamaV2Config,10    ExLlamaV2Tokenizer11)12from exllamav2.generator import ExLlamaV2Sampler, ExLlamaV2StreamingGenerator13 14from modules import shared15from modules.logging_colors import logger16from modules.text_generation import get_max_prompt_length17 18try:19    import flash_attn20except ModuleNotFoundError:21    logger.warning(22        'You are running ExLlamaV2 without flash-attention. This will cause the VRAM usage '23        'to be a lot higher than it could be.\n'24        'Try installing flash-attention following the instructions here: '25        'https://github.com/Dao-AILab/flash-attention#installation-and-features'26    )27    pass28except Exception:29    logger.warning('Failed to load flash-attention due to the following error:\n')30    traceback.print_exc()31 32 33class Exllamav2Model:34    def __init__(self):35        pass36 37    @classmethod38    def from_pretrained(self, path_to_model):39 40        path_to_model = Path(f'{shared.args.model_dir}') / Path(path_to_model)41 42        config = ExLlamaV2Config()43        config.model_dir = str(path_to_model)44        config.prepare()45 46        config.max_seq_len = shared.args.max_seq_len47        config.scale_pos_emb = shared.args.compress_pos_emb48        config.scale_alpha_value = shared.args.alpha_value49        config.no_flash_attn = shared.args.no_flash_attn50        config.num_experts_per_token = int(shared.args.num_experts_per_token)51 52        model = ExLlamaV2(config)53 54        split = None55        if shared.args.gpu_split:56            split = [float(alloc) for alloc in shared.args.gpu_split.split(",")]57 58        model.load(split)59 60        tokenizer = ExLlamaV2Tokenizer(config)61        if shared.args.cache_8bit:62            cache = ExLlamaV2Cache_8bit(model)63        else:64            cache = ExLlamaV2Cache(model)65 66        generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer)67 68        result = self()69        result.model = model70        result.cache = cache71        result.tokenizer = tokenizer72        result.generator = generator73        result.loras = None74        return result, result75 76    def encode(self, string, **kwargs):77        return self.tokenizer.encode(string, add_bos=True, encode_special_tokens=True)78 79    def decode(self, ids, **kwargs):80        if isinstance(ids, list):81            ids = torch.tensor([ids])82        elif isinstance(ids, torch.Tensor) and ids.numel() == 1:83            ids = ids.view(1, -1)84 85        return self.tokenizer.decode(ids, decode_special_tokens=True)[0]86 87    def get_logits(self, token_ids, **kwargs):88        self.cache.current_seq_len = 089        if token_ids.shape[-1] > 1:90            self.model.forward(token_ids[:, :-1], self.cache, input_mask=None, preprocess_only=True, loras=self.loras)91 92        return self.model.forward(token_ids[:, -1:], self.cache, input_mask=None, loras=self.loras, **kwargs).float().cpu()93 94    def generate_with_streaming(self, prompt, state):95        settings = ExLlamaV2Sampler.Settings()96        settings.temperature = state['temperature']97        settings.top_k = state['top_k']98        settings.top_p = state['top_p']99        settings.min_p = state['min_p']100        settings.tfs = state['tfs']101        settings.typical = state['typical_p']102        settings.mirostat = state['mirostat_mode'] == 2103        settings.mirostat_tau = state['mirostat_tau']104        settings.mirostat_eta = state['mirostat_eta']105        settings.token_repetition_penalty = state['repetition_penalty']106        settings.token_repetition_range = -1 if state['repetition_penalty_range'] <= 0 else state['repetition_penalty_range']107        if state['ban_eos_token']:108            settings.disallow_tokens(self.tokenizer, [self.tokenizer.eos_token_id])109 110        if state['custom_token_bans']:111            to_ban = [int(x) for x in state['custom_token_bans'].split(',')]112            if len(to_ban) > 0:113                settings.disallow_tokens(self.tokenizer, to_ban)114 115        ids = self.tokenizer.encode(prompt, add_bos=state['add_bos_token'], encode_special_tokens=True)116        ids = ids[:, -get_max_prompt_length(state):]117 118        if state['auto_max_new_tokens']:119            max_new_tokens = state['truncation_length'] - ids.shape[-1]120        else:121            max_new_tokens = state['max_new_tokens']122 123        self.generator.begin_stream(ids, settings, loras=self.loras)124 125        decoded_text = ''126        for i in range(max_new_tokens):127            chunk, eos, _ = self.generator.stream()128            if eos or shared.stop_everything:129                break130 131            decoded_text += chunk132            yield decoded_text133 134    def generate(self, prompt, state):135        output = ''136        for output in self.generate_with_streaming(prompt, state):137            pass138 139        return output140