forestcalled/text-generation-webui

1import os2from pathlib import Path3from typing import Any, Dict, Optional, Union4 5import torch6from torch.nn import CrossEntropyLoss7from transformers import GenerationConfig, PretrainedConfig, PreTrainedModel8from transformers.modeling_outputs import CausalLMOutputWithPast9 10from modules import shared11from modules.logging_colors import logger12 13try:14    from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig15except:16    logger.warning('Exllama module failed to load. Will attempt to load from repositories.')17    try:18        from modules.relative_imports import RelativeImport19 20        with RelativeImport("repositories/exllama"):21            from model import ExLlama, ExLlamaCache, ExLlamaConfig22    except:23        logger.error("Could not find repositories/exllama/. Make sure that exllama is cloned inside repositories/ and is up to date.")24        raise25 26 27class ExllamaHF(PreTrainedModel):28    def __init__(self, config: ExLlamaConfig):29        super().__init__(PretrainedConfig())30        self.ex_config = config31        self.ex_model = ExLlama(self.ex_config)32        self.generation_config = GenerationConfig()33        self.lora = None34 35        self.ex_cache = ExLlamaCache(self.ex_model)36        self.past_seq = None37 38        if shared.args.cfg_cache:39            self.ex_cache_negative = ExLlamaCache(self.ex_model)40            self.past_seq_negative = None41 42    def _validate_model_class(self):43        pass44 45    def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):46        pass47 48    def prepare_inputs_for_generation(self, input_ids, **kwargs):49        return {'input_ids': input_ids, **kwargs}50 51    @property52    def device(self) -> torch.device:53        return torch.device(0)54 55    def __call__(self, *args, **kwargs):56        use_cache = kwargs.get('use_cache', True)57        labels = kwargs.get('labels', None)58        past_key_values = kwargs.get('past_key_values', None)59 60        if len(args) > 0:61            if not shared.args.cfg_cache:62                logger.error("Please enable the cfg-cache option to use CFG with ExLlama_HF.")63                return64 65            input_ids = args[0]66            is_negative = True67            past_seq = self.past_seq_negative68            ex_cache = self.ex_cache_negative69        else:70            input_ids = kwargs['input_ids']71            is_negative = False72            past_seq = self.past_seq73            ex_cache = self.ex_cache74 75        seq = input_ids[0].tolist()76        if is_negative and past_key_values is not None:77            seq = past_key_values + seq78 79        seq_tensor = torch.tensor(seq)80        reset = True81 82        # Make the forward call83        if labels is None:84            if past_seq is not None:85                min_length = min(past_seq.shape[0], seq_tensor.shape[0])86                indices = torch.nonzero(~torch.eq(past_seq[:min_length], seq_tensor[:min_length]))87                if len(indices) > 0:88                    longest_prefix = indices[0].item()89                else:90                    longest_prefix = min_length91 92                if longest_prefix > 0:93                    reset = False94                    ex_cache.current_seq_len = longest_prefix95                    if len(seq_tensor) - longest_prefix > 1:96                        self.ex_model.forward(seq_tensor[longest_prefix:-1].view(1, -1), ex_cache, preprocess_only=True, lora=self.lora)97                    elif len(seq_tensor) == longest_prefix:98                        # Very tricky: if the prefix we are reusing *is* the input_ids, then we have to back up the cache pointer by one,99                        # because we feed input_ids[-1] to forward() below, but that last token is already in the cache!100                        ex_cache.current_seq_len -= 1101 102            if reset:103                ex_cache.current_seq_len = 0104                if len(seq_tensor) > 1:105                    self.ex_model.forward(seq_tensor[:-1].view(1, -1), ex_cache, preprocess_only=True, lora=self.lora)106 107            logits = self.ex_model.forward(seq_tensor[-1:].view(1, -1), ex_cache, lora=self.lora).to(input_ids.device)108        else:109            ex_cache.current_seq_len = 0110            logits = self.ex_model.forward(seq_tensor.view(1, -1), ex_cache, last_id_only=False, lora=self.lora)111 112        if is_negative:113            self.past_seq_negative = seq_tensor114        else:115            self.past_seq = seq_tensor116 117        loss = None118        if labels is not None:119            # Shift so that tokens < n predict n120            shift_logits = logits[..., :-1, :].contiguous()121            shift_labels = labels[..., 1:].contiguous()122            # Flatten the tokens123            loss_fct = CrossEntropyLoss()124            shift_logits = shift_logits.view(-1, logits.shape[-1])125            shift_labels = shift_labels.view(-1)126            # Enable model parallelism127            shift_labels = shift_labels.to(shift_logits.device)128            loss = loss_fct(shift_logits, shift_labels)129 130        return CausalLMOutputWithPast(logits=logits, past_key_values=seq if use_cache else None, loss=loss)131 132    @classmethod133    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):134        assert len(model_args) == 0 and len(kwargs) == 0, "extra args is currently not supported"135        if isinstance(pretrained_model_name_or_path, str):136            pretrained_model_name_or_path = Path(pretrained_model_name_or_path)137 138        pretrained_model_name_or_path = Path(f'{shared.args.model_dir}') / Path(pretrained_model_name_or_path)139        config = ExLlamaConfig(pretrained_model_name_or_path / 'config.json')140 141        # from 'oobabooga/text-generation-webui/modules/exllama.py'142        weight_path = None143        for ext in ['.safetensors', '.pt', '.bin']:144            found = list(pretrained_model_name_or_path.glob(f"*{ext}"))145            if len(found) > 0:146                weight_path = found[-1]147                break148        assert weight_path is not None, f'could not find weight in "{pretrained_model_name_or_path}"'149 150        config.model_path = str(weight_path)151        config.max_seq_len = shared.args.max_seq_len152        config.compress_pos_emb = shared.args.compress_pos_emb153        if shared.args.gpu_split:154            config.set_auto_map(shared.args.gpu_split)155            config.gpu_peer_fix = True156 157        if shared.args.alpha_value > 1 and shared.args.rope_freq_base == 0:158            config.alpha_value = shared.args.alpha_value159            config.calculate_rotary_embedding_base()160        elif shared.args.rope_freq_base > 0:161            config.rotary_embedding_base = shared.args.rope_freq_base162 163        if torch.version.hip:164            config.rmsnorm_no_half2 = True165            config.rope_no_half2 = True166            config.matmul_no_half2 = True167            config.silu_no_half2 = True168 169        # This slowes down a bit but align better with autogptq generation.170        # TODO: Should give user choice to tune the exllama config171        # config.fused_attn = False172        # config.fused_mlp_thd = 0173 174        return ExllamaHF(config)175