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 RoPE, shared11from modules.logging_colors import logger12 13try:14    import llama_cpp15except:16    llama_cpp = None17 18try:19    import llama_cpp_cuda20except:21    llama_cpp_cuda = None22 23try:24    import llama_cpp_cuda_tensorcores25except:26    llama_cpp_cuda_tensorcores = None27 28 29def llama_cpp_lib():30    if shared.args.cpu and llama_cpp is not None:31        return llama_cpp32    elif shared.args.tensorcores and llama_cpp_cuda_tensorcores is not None:33        return llama_cpp_cuda_tensorcores34    elif llama_cpp_cuda is not None:35        return llama_cpp_cuda36    else:37        return llama_cpp38 39 40class LlamacppHF(PreTrainedModel):41    def __init__(self, model, path):42        super().__init__(PretrainedConfig())43        self.model = model44        self.generation_config = GenerationConfig()45 46        self.past_seq = None47        self.llamacpp_cache = {48            'n_tokens': self.model.n_tokens,49            'input_ids': self.model.input_ids,50            'scores': self.model.scores,51            'ctx': self.model._ctx52        }53 54        if shared.args.cfg_cache:55            self.past_seq_negative = None56            self.llamacpp_cache_negative = {57                'n_tokens': self.model.n_tokens,58                'input_ids': self.model.input_ids.copy(),59                'scores': self.model.scores.copy(),60                'ctx': llama_cpp_lib().llama_new_context_with_model(model.model, model.context_params)61            }62 63    def _validate_model_class(self):64        pass65 66    def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):67        pass68 69    def prepare_inputs_for_generation(self, input_ids, **kwargs):70        return {'input_ids': input_ids, **kwargs}71 72    def save_cache(self):73        self.llamacpp_cache.update({74            'n_tokens': self.model.n_tokens,75            'input_ids': self.model.input_ids,76            'scores': self.model.scores,77            'ctx': self.model._ctx78        })79 80    def save_negative_cache(self):81        self.llamacpp_cache_negative.update({82            'n_tokens': self.model.n_tokens,83            'input_ids': self.model.input_ids,84            'scores': self.model.scores,85            'ctx': self.model._ctx86        })87 88    def load_cache(self):89        self.model.n_tokens = self.llamacpp_cache['n_tokens']90        self.model.input_ids = self.llamacpp_cache['input_ids']91        self.model.scores = self.llamacpp_cache['scores']92        self.model._ctx = self.llamacpp_cache['ctx']93 94    def load_negative_cache(self):95        self.model.n_tokens = self.llamacpp_cache_negative['n_tokens']96        self.model.input_ids = self.llamacpp_cache_negative['input_ids']97        self.model.scores = self.llamacpp_cache_negative['scores']98        self.model._ctx = self.llamacpp_cache_negative['ctx']99 100    @property101    def device(self) -> torch.device:102        return torch.device(0)103 104    def __call__(self, *args, **kwargs):105        use_cache = kwargs.get('use_cache', True)106        labels = kwargs.get('labels', None)107        past_key_values = kwargs.get('past_key_values', None)108 109        if len(args) > 0:110            if not shared.args.cfg_cache:111                logger.error("Please enable the cfg-cache option to use CFG with llamacpp_HF.")112                return113 114            input_ids = args[0]115            is_negative = True116            past_seq = self.past_seq_negative117            self.load_negative_cache()118        else:119            input_ids = kwargs['input_ids']120            is_negative = False121            past_seq = self.past_seq122            self.load_cache()123 124        seq = input_ids[0].tolist()125        if is_negative and past_key_values is not None:126            seq = past_key_values + seq127 128        seq_tensor = torch.tensor(seq)129        reset = True130 131        # Make the forward call. The prefix-match code has been adapted from132        # https://github.com/abetlen/llama-cpp-python/commit/f4090a0bb2a2a25acfe28d31c82cc1aa273bedee133        if labels is None:134            if past_seq is not None:135                min_length = min(past_seq.shape[0], seq_tensor.shape[0])136                indices = torch.nonzero(~torch.eq(past_seq[:min_length], seq_tensor[:min_length]))137                if len(indices) > 0:138                    longest_prefix = indices[0].item()139                else:140                    longest_prefix = min_length141 142                if longest_prefix > 0:143                    reset = False144                    self.model.n_tokens = longest_prefix145                    if len(seq_tensor) - longest_prefix > 0:146                        self.model.eval(seq[longest_prefix:])147 148            if reset:149                self.model.reset()150                self.model.eval(seq)151 152            logits = torch.tensor(self.model.scores[self.model.n_tokens - 1, :]).view(1, 1, -1).to(input_ids.device)153        else:154            self.model.reset()155            self.model.eval(seq)156            logits = torch.tensor(self.model.eval_logits)157            logits = logits.view(1, logits.shape[0], logits.shape[1]).to(input_ids.device)158 159        if is_negative:160            self.save_negative_cache()161            self.past_seq_negative = seq_tensor162        else:163            self.save_cache()164            self.past_seq = seq_tensor165 166        loss = None167        if labels is not None:168            # Shift so that tokens < n predict n169            shift_logits = logits[..., :-1, :].contiguous()170            shift_labels = labels[..., 1:].contiguous()171            # Flatten the tokens172            loss_fct = CrossEntropyLoss()173            shift_logits = shift_logits.view(-1, logits.shape[-1])174            shift_labels = shift_labels.view(-1)175            # Enable model parallelism176            shift_labels = shift_labels.to(shift_logits.device)177            loss = loss_fct(shift_logits, shift_labels)178 179        return CausalLMOutputWithPast(logits=logits, past_key_values=seq if use_cache else None, loss=loss)180 181    @classmethod182    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):183        assert len(model_args) == 0 and len(kwargs) == 0, "extra args is currently not supported"184 185        if isinstance(pretrained_model_name_or_path, str):186            pretrained_model_name_or_path = Path(pretrained_model_name_or_path)187 188        path = Path(f'{shared.args.model_dir}') / Path(pretrained_model_name_or_path)189        if path.is_file():190            model_file = path191        else:192            model_file = list(path.glob('*.gguf'))[0]193 194        logger.info(f"llama.cpp weights detected: {model_file}\n")195 196        if shared.args.tensor_split is None or shared.args.tensor_split.strip() == '':197            tensor_split_list = None198        else:199            tensor_split_list = [float(x) for x in shared.args.tensor_split.strip().split(",")]200 201        params = {202            'model_path': str(model_file),203            'n_ctx': shared.args.n_ctx,204            'n_threads': shared.args.threads or None,205            'n_threads_batch': shared.args.threads_batch or None,206            'n_batch': shared.args.n_batch,207            'use_mmap': not shared.args.no_mmap,208            'use_mlock': shared.args.mlock,209            'mul_mat_q': not shared.args.no_mul_mat_q,210            'numa': shared.args.numa,211            'n_gpu_layers': shared.args.n_gpu_layers,212            'rope_freq_base': RoPE.get_rope_freq_base(shared.args.alpha_value, shared.args.rope_freq_base),213            'tensor_split': tensor_split_list,214            'rope_freq_scale': 1.0 / shared.args.compress_pos_emb,215            'logits_all': shared.args.logits_all,216            'offload_kqv': not shared.args.no_offload_kqv217        }218 219        Llama = llama_cpp_lib().Llama220        model = Llama(**params)221 222        return LlamacppHF(model, model_file)223