forestcalled/text-generation-webui

1from pathlib import Path2 3import torch4import torch.nn.functional as F5from torch import version as torch_version6 7from modules import shared8from modules.logging_colors import logger9from modules.models import clear_torch_cache10from modules.text_generation import get_max_prompt_length11 12try:13    from exllama.generator import ExLlamaGenerator14    from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig15    from exllama.tokenizer import ExLlamaTokenizer16except:17    logger.warning('exllama module failed to import. Will attempt to import from repositories/.')18    try:19        from modules.relative_imports import RelativeImport20 21        with RelativeImport("repositories/exllama"):22            from generator import ExLlamaGenerator23            from model import ExLlama, ExLlamaCache, ExLlamaConfig24            from tokenizer import ExLlamaTokenizer25    except:26        logger.error(27            "Could not find repositories/exllama. Please ensure that exllama"28            " (https://github.com/turboderp/exllama) is cloned inside repositories/ and is up to date."29        )30        raise31 32 33class ExllamaModel: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        tokenizer_model_path = path_to_model / "tokenizer.model"42        model_config_path = path_to_model / "config.json"43 44        # Find the model checkpoint45        model_path = None46        for ext in ['.safetensors', '.pt', '.bin']:47            found = list(path_to_model.glob(f"*{ext}"))48            if len(found) > 0:49                if len(found) > 1:50                    logger.warning(f'More than one {ext} model has been found. The last one will be selected. It could be wrong.')51 52                model_path = found[-1]53                break54 55        config = ExLlamaConfig(str(model_config_path))56        config.model_path = str(model_path)57        config.max_seq_len = shared.args.max_seq_len58        config.compress_pos_emb = shared.args.compress_pos_emb59        if shared.args.gpu_split:60            config.set_auto_map(shared.args.gpu_split)61            config.gpu_peer_fix = True62 63        if shared.args.alpha_value > 1 and shared.args.rope_freq_base == 0:64            config.alpha_value = shared.args.alpha_value65            config.calculate_rotary_embedding_base()66        elif shared.args.rope_freq_base > 0:67            config.rotary_embedding_base = shared.args.rope_freq_base68 69        if torch_version.hip:70            config.rmsnorm_no_half2 = True71            config.rope_no_half2 = True72            config.matmul_no_half2 = True73            config.silu_no_half2 = True74 75        model = ExLlama(config)76        tokenizer = ExLlamaTokenizer(str(tokenizer_model_path))77        cache = ExLlamaCache(model)78        generator = ExLlamaGenerator(model, tokenizer, cache)79 80        result = self()81        result.config = config82        result.model = model83        result.cache = cache84        result.tokenizer = tokenizer85        result.generator = generator86        return result, result87 88    def encode(self, string, **kwargs):89        return self.tokenizer.encode(string, max_seq_len=self.model.config.max_seq_len, add_bos=True)90 91    def decode(self, ids, **kwargs):92        if isinstance(ids, list):93            ids = torch.tensor([ids])94        elif isinstance(ids, torch.Tensor) and ids.numel() == 1:95            ids = ids.view(1, -1)96 97        return self.tokenizer.decode(ids)[0]98 99    def get_logits(self, token_ids, **kwargs):100        self.cache.current_seq_len = 0101        if token_ids.shape[-1] > 1:102            self.model.forward(token_ids[:, :-1], self.cache, input_mask=None, preprocess_only=True)103 104        return self.model.forward(token_ids[:, -1:], self.cache, **kwargs).float().cpu()105 106    def generate_with_streaming(self, prompt, state):107 108        # The cache batch size must be 2 for CFG and 1 otherwise109        if state['guidance_scale'] == 1:110            if self.cache.batch_size == 2:111                del self.cache112                clear_torch_cache()113                self.cache = ExLlamaCache(self.model)114                self.generator = ExLlamaGenerator(self.model, self.tokenizer, self.cache)115        else:116            if self.cache.batch_size == 1:117                del self.cache118                clear_torch_cache()119                self.cache = ExLlamaCache(self.model, batch_size=2)120                self.generator = ExLlamaGenerator(self.model, self.tokenizer, self.cache)121 122        self.generator.settings.temperature = state['temperature']123        self.generator.settings.top_p = state['top_p']124        self.generator.settings.top_k = state['top_k']125        self.generator.settings.typical = state['typical_p']126        self.generator.settings.token_repetition_penalty_max = state['repetition_penalty']127        self.generator.settings.token_repetition_penalty_sustain = -1 if state['repetition_penalty_range'] <= 0 else state['repetition_penalty_range']128        if state['ban_eos_token']:129            self.generator.disallow_tokens([self.tokenizer.eos_token_id])130        else:131            self.generator.disallow_tokens(None)132 133        if state['custom_token_bans']:134            to_ban = [int(x) for x in state['custom_token_bans'].split(',')]135            if len(to_ban) > 0:136                self.generator.disallow_tokens(to_ban)137 138        # Case 1: no CFG139        if state['guidance_scale'] == 1:140            self.generator.end_beam_search()141 142            # Tokenizing the input143            ids = self.generator.tokenizer.encode(prompt, max_seq_len=self.model.config.max_seq_len)144            if state['add_bos_token']:145                ids = torch.cat(146                    [torch.tensor([[self.tokenizer.bos_token_id]]).to(ids.device),147                     ids], dim=1148                ).to(torch.int64)149            ids = ids[:, -get_max_prompt_length(state):]150            if state['auto_max_new_tokens']:151                max_new_tokens = state['truncation_length'] - ids.shape[-1]152            else:153                max_new_tokens = state['max_new_tokens']154 155            self.generator.gen_begin_reuse(ids)156            initial_len = self.generator.sequence[0].shape[0]157            has_leading_space = False158 159            for i in range(max_new_tokens):160                token = self.generator.gen_single_token()161                if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('▁'):162                    has_leading_space = True163 164                decoded_text = self.generator.tokenizer.decode(self.generator.sequence[0][initial_len:])165                if has_leading_space:166                    decoded_text = ' ' + decoded_text167 168                # Check the partial unicode character169                if chr(0xfffd) in decoded_text:170                    is_last = i == max_new_tokens - 1171                    is_stopping = token.item() == self.generator.tokenizer.eos_token_id or shared.stop_everything172                    # If we are not at the end of the generation, we skip this token173                    if not (is_last or is_stopping):174                        continue175 176                if token.item() == self.generator.tokenizer.eos_token_id or shared.stop_everything:177                    break178 179                yield decoded_text180 181        # Case 2: CFG182        # Copied from https://github.com/turboderp/exllama/blob/master/example_cfg.py183        else:184            alpha = state['guidance_scale']185            prompts = [prompt, state['negative_prompt'] or '']186 187            ids, mask = self.tokenizer.encode(188                prompts,189                return_mask=True,190                max_seq_len=self.model.config.max_seq_len,191                add_bos=state['add_bos_token']192            )193            if state['auto_max_new_tokens']:194                max_new_tokens = state['truncation_length'] - ids[0].shape[-1]195            else:196                max_new_tokens = state['max_new_tokens']197 198            self.generator.gen_begin(ids, mask=mask)199            initial_len = self.generator.sequence[0].shape[0]200            has_leading_space = False201 202            for i in range(max_new_tokens):203                logits = self.model.forward(self.generator.sequence[:, -1:], self.cache, input_mask=mask)204                self.generator.apply_rep_penalty(logits)205 206                logits = F.log_softmax(logits, dim=-1)207                logits_mixed = alpha * logits[0] + (1 - alpha) * logits[1]208 209                token, _ = self.generator.sample_current(logits_mixed)210                if i == 0 and self.generator.tokenizer.tokenizer.IdToPiece(int(token)).startswith('▁'):211                    has_leading_space = True212 213                decoded_text = self.generator.tokenizer.decode(self.generator.sequence[0][initial_len:])214                if has_leading_space:215                    decoded_text = ' ' + decoded_text216 217                # Check the partial unicode character218                if chr(0xfffd) in decoded_text:219                    is_last = i == max_new_tokens - 1220                    is_stopping = token.item() == self.tokenizer.eos_token_id or shared.stop_everything221                    # If we are not at the end of the generation, we skip this token222                    if not (is_last or is_stopping):223                        continue224 225                yield decoded_text226                if token.item() == self.tokenizer.eos_token_id or shared.stop_everything:227                    break228 229                batch_token = token.repeat(2, 1)230                self.generator.gen_accept_token(batch_token)231 232    def generate(self, prompt, state):233        output = ''234        for output in self.generate_with_streaming(prompt, state):235            pass236 237        return output238