forestcalled/text-generation-webui

1import math2 3import torch4import transformers5from transformers import LogitsWarper, is_torch_xpu_available6from transformers.generation.logits_process import (7    LogitNormalization,8    LogitsProcessor,9    LogitsProcessorList,10    TemperatureLogitsWarper11)12 13global_scores = None14 15 16class MinPLogitsWarper(LogitsWarper):17    def __init__(self, min_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):18        if min_p < 0 or min_p > 1.0:19            raise ValueError(f"`min_p` has to be a float >= 0 and <= 1, but is {min_p}")20        self.min_p = min_p21        self.filter_value = filter_value22        self.min_tokens_to_keep = min_tokens_to_keep23 24    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:25        # Convert logits to probabilities26        probs = torch.softmax(scores, dim=-1)27        # Get the probability of the top token for each sequence in the batch28        top_probs, _ = probs.max(dim=-1, keepdim=True)29        # Calculate the actual min_p threshold by scaling min_p with the top token's probability30        scaled_min_p = self.min_p * top_probs31        # Create a mask for tokens that have a probability less than the scaled min_p32        tokens_to_remove = probs < scaled_min_p33 34        sorted_indices = torch.argsort(scores, descending=True, dim=-1)35        sorted_indices_to_remove = torch.gather(tokens_to_remove, dim=-1, index=sorted_indices)36 37        if self.min_tokens_to_keep > 1:38            # Keep at least min_tokens_to_keep39            sorted_indices_to_remove[..., : self.min_tokens_to_keep] = False40 41        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)42        scores = scores.masked_fill(indices_to_remove, self.filter_value)43        return scores44 45 46class TailFreeLogitsWarper(LogitsWarper):47    def __init__(self, tfs: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):48        tfs = float(tfs)49        if tfs < 0 or tfs > 1.0:50            raise ValueError(f"`tfs` has to be a float >= 0 and <= 1, but is {tfs}")51        self.tfs = tfs52        self.filter_value = filter_value53        self.min_tokens_to_keep = min_tokens_to_keep54 55    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:56        sorted_logits, sorted_indices = torch.sort(scores, descending=True)57        probs = sorted_logits.softmax(dim=-1)58 59        # Compute second derivative normalized CDF60        d2 = probs.diff().diff().abs()61        normalized_d2 = d2 / d2.sum(dim=-1, keepdim=True)62        normalized_d2_cdf = normalized_d2.cumsum(dim=-1)63 64        # Remove tokens with CDF value above the threshold (token with 0 are kept)65        sorted_indices_to_remove = normalized_d2_cdf > self.tfs66 67        # Centre the distribution around the cutoff as in the original implementation of the algorithm68        sorted_indices_to_remove = torch.cat(69            (70                torch.zeros(scores.shape[0], 1, dtype=torch.bool, device=scores.device),71                sorted_indices_to_remove,72                torch.ones(scores.shape[0], 1, dtype=torch.bool, device=scores.device),73            ),74            dim=-1,75        )76 77        if self.min_tokens_to_keep > 1:78            # Keep at least min_tokens_to_keep79            sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 080 81        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)82        scores = scores.masked_fill(indices_to_remove, self.filter_value)83        return scores84 85 86class TopALogitsWarper(LogitsWarper):87    def __init__(self, top_a: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):88        top_a = float(top_a)89        if top_a < 0 or top_a > 1.0:90            raise ValueError(f"`top_a` has to be a float >= 0 and <= 1, but is {top_a}")91        self.top_a = top_a92        self.filter_value = filter_value93        self.min_tokens_to_keep = min_tokens_to_keep94 95    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:96        sorted_logits, sorted_indices = torch.sort(scores, descending=True)97        probs = sorted_logits.softmax(dim=-1)98 99        # Remove tokens with probability less than top_a*(max(probs))^2 (token with 0 are kept)100        probs_max = probs[..., 0, None]101        sorted_indices_to_remove = probs < probs_max * probs_max * self.top_a102 103        if self.min_tokens_to_keep > 1:104            # Keep at least min_tokens_to_keep105            sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 0106 107        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)108        scores = scores.masked_fill(indices_to_remove, self.filter_value)109        return scores110 111 112class MirostatLogitsWarper(LogitsWarper):113    def __init__(self, mirostat_mode: int, mirostat_tau: float, mirostat_eta: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):114        if mirostat_mode not in [2]:115            raise ValueError(f"`mirostat` has to be a an integer 2, but is {mirostat_mode}")116        self.mirostat_mode = mirostat_mode117        self.mirostat_eta = mirostat_eta118        self.mirostat_tau = mirostat_tau119        self.filter_value = filter_value120        self.min_tokens_to_keep = min_tokens_to_keep121        self.mu = 2 * self.mirostat_tau122        self.e = 0123 124    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:125        logits = scores[0]126        sorted_logits, sorted_indices = torch.sort(logits, descending=True)127        prob_original = torch.softmax(sorted_logits, dim=-1).tolist()  # candidates128 129        # Truncate the words with surprise values greater than mu130        for i, candidate in enumerate(prob_original):131            if candidate > 0 and -math.log2(candidate) > self.mu:132                if (i == 0):133                    sorted_logits = sorted_logits[:1]134                else:135                    sorted_logits = sorted_logits[:i]136                break137 138        # Normalize the probabilities of the remaining words139        if is_torch_xpu_available():140            prob_topk = torch.softmax(sorted_logits, dim=0).to("xpu")141            prev_i = torch.multinomial(prob_topk, num_samples=1, replacement=True).to("xpu")142        else:143            prob_topk = torch.softmax(sorted_logits, dim=0).to('cuda')144            prev_i = torch.multinomial(prob_topk, num_samples=1, replacement=True).to('cuda')145 146        observed_surprise = -math.log2(prob_topk[prev_i])147        self.e = observed_surprise - self.mirostat_tau148 149        # Update mu using the learning rate and error150        self.mu -= self.mirostat_eta * self.e151 152        sorted_indices_to_remove = torch.ones_like(scores[0], dtype=torch.bool)153        sorted_indices_to_remove[prev_i] = False154 155        indices_to_remove = sorted_indices_to_remove.unsqueeze(0).scatter(1, sorted_indices.unsqueeze(0), sorted_indices_to_remove.unsqueeze(0))156        scores = scores.masked_fill(indices_to_remove, self.filter_value)157        return scores158 159 160class SpyLogitsWarper(LogitsWarper):161    def __init__(self):162        pass163 164    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:165        global global_scores166        global_scores = scores167        return scores168 169 170class RepetitionPenaltyLogitsProcessorWithRange(LogitsProcessor):171    '''172    Copied from the transformers library173    '''174 175    def __init__(self, penalty: float, presence_penalty: float, frequency_penalty: float, _range: int):176        if not (penalty > 0):177            raise ValueError(f"`penalty` has to be strictly positive, but is {penalty}")178 179        self.penalty = penalty180        self.presence_penalty = presence_penalty181        self.frequency_penalty = frequency_penalty182        self._range = _range183 184    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:185        input_ids = input_ids[:, -self._range:]186 187        # We loop here because torch.unique() needs to process each row separately in the188        # case that batch_size > 1.189        for input_ids_row, scores_row in zip(input_ids, scores):190            unique_ids, counts = torch.unique(input_ids_row, return_counts=True)191            score = torch.gather(scores_row, 0, unique_ids)192 193            # multiplicative repetition penalty194            # if score < 0 then repetition penalty has to be multiplied to reduce the previous token probability195            score = torch.where(score < 0, score * self.penalty, score / self.penalty)196            scores_row.scatter_(0, unique_ids, score)197 198            # presence_penalty and frequency_penalty199            raw_presence_penalty = (counts > 0).to(scores.dtype)200            raw_frequency_penalty = counts.to(scores.dtype)201            additive_penalty = raw_presence_penalty*self.presence_penalty + raw_frequency_penalty*self.frequency_penalty202            scores_row.scatter_add_(0, unique_ids, -additive_penalty)203 204        return scores205 206 207def get_logits_warper_patch(self, generation_config):208    warpers = self._get_logits_warper_old(generation_config)209    warpers_to_add = LogitsProcessorList()210    min_tokens_to_keep = 2 if generation_config.num_beams > 1 else 1211 212    if generation_config.mirostat_mode is not None and generation_config.mirostat_mode == 2:213        warpers_to_add.append(MirostatLogitsWarper(mirostat_mode=generation_config.mirostat_mode, mirostat_eta=generation_config.mirostat_eta, mirostat_tau=generation_config.mirostat_tau, min_tokens_to_keep=min_tokens_to_keep))214        # We need to disable samplers other than temperature215        for warper in warpers:216            if not isinstance(warper, TemperatureLogitsWarper):217                warpers.remove(warper)218    else:219        if generation_config.tfs is not None and 0.0 <= generation_config.tfs < 1.0:220            warpers_to_add.append(TailFreeLogitsWarper(tfs=generation_config.tfs, min_tokens_to_keep=min_tokens_to_keep))221        if generation_config.top_a is not None and 0.0 < generation_config.top_a <= 1.0:222            warpers_to_add.append(TopALogitsWarper(top_a=generation_config.top_a, min_tokens_to_keep=min_tokens_to_keep))223        if generation_config.min_p is not None and 0.0 < generation_config.min_p <= 1.0:224            warpers_to_add.append(MinPLogitsWarper(min_p=generation_config.min_p, min_tokens_to_keep=min_tokens_to_keep))225 226    if len(warpers) > 0 and isinstance(warpers[-1], LogitNormalization):227        normalize = warpers.pop(-1)228    else:229        normalize = None230 231    warpers += warpers_to_add232    if generation_config.temperature_last:233        temperature_idx = None234        for i in range(len(warpers)):235            if warpers[i].__class__.__name__ == 'TemperatureLogitsWarper':236                temperature_idx = i237                break238 239        if temperature_idx is not None:240            warpers = warpers[:temperature_idx] + warpers[temperature_idx + 1:] + [warpers[temperature_idx]]241            warpers = LogitsProcessorList(warpers)242 243    if normalize is not None:244        warpers.append(normalize)245 246    warpers.append(SpyLogitsWarper())247    # for i in range(len(warpers)):248    #     print(warpers[i].__class__.__name__)249    return warpers250 251 252def get_logits_processor_patch(self, **kwargs):253    repetition_penalty = kwargs['generation_config'].repetition_penalty254    presence_penalty = kwargs['generation_config'].presence_penalty255    frequency_penalty = kwargs['generation_config'].frequency_penalty256    repetition_penalty_range = kwargs['generation_config'].repetition_penalty_range257    do_rep_pen_hijack = (repetition_penalty > 1) or (presence_penalty != 0) or (frequency_penalty != 0)258    if do_rep_pen_hijack:259        # Make sure that a RepetitionPenaltyLogitsProcessor will be created260        kwargs['generation_config'].repetition_penalty = 1.1  # must set to some value > 1261 262    result = self._get_logits_processor_old(**kwargs)263 264    if do_rep_pen_hijack:265        for i in range(len(result)):266            if result[i].__class__.__name__ == 'RepetitionPenaltyLogitsProcessor':267                result[i] = RepetitionPenaltyLogitsProcessorWithRange(repetition_penalty, presence_penalty, frequency_penalty, repetition_penalty_range)268 269    return result270 271 272def generation_config_init_patch(self, **kwargs):273    self.__init___old(**kwargs)274    self.min_p = kwargs.pop("min_p", 0.0)275    self.tfs = kwargs.pop("tfs", 1.0)276    self.top_a = kwargs.pop("top_a", 0.0)277    self.mirostat_mode = kwargs.pop("mirostat_mode", 0)278    self.mirostat_eta = kwargs.pop("mirostat_eta", 0.1)279    self.mirostat_tau = kwargs.pop("mirostat_tau", 5)280    self.repetition_penalty_range = kwargs.pop("repetition_penalty_range", 0)281    self.presence_penalty = kwargs.pop("presence_penalty", 0)282    self.frequency_penalty = kwargs.pop("frequency_penalty", 0)283    self.temperature_last = kwargs.pop("temperature_last", False)284 285 286def hijack_samplers():287    transformers.GenerationMixin._get_logits_warper_old = transformers.GenerationMixin._get_logits_warper288    transformers.GenerationMixin._get_logits_warper = get_logits_warper_patch289 290    transformers.GenerationMixin._get_logits_processor_old = transformers.GenerationMixin._get_logits_processor291    transformers.GenerationMixin._get_logits_processor = get_logits_processor_patch292 293    transformers.GenerationConfig.__init___old = transformers.GenerationConfig.__init__294    transformers.GenerationConfig.__init__ = generation_config_init_patch295