forestcalled/text-generation-webui

1'''2This file has been 100% copied from this PR to the Transformers library:3https://github.com/huggingface/transformers/pull/275574 5Author: Saibo-creator6Author GitHub: https://github.com/Saibo-creator7 8All credits go to the author.9'''10 11import math12 13import torch14from transformers.generation.logits_process import LogitsProcessor15from transformers.utils import add_start_docstrings16 17LOGITS_PROCESSOR_INPUTS_DOCSTRING = r"""18    Args:19        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):20            Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)21        scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):22            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam23            search or log softmax for each vocabulary token when using beam search24 25    Return:26        `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.27 28"""29 30 31class GrammarConstrainedLogitsProcessor(LogitsProcessor):32    def __init__(self, grammar_constraint):33        self.last_size = None34        self.grammar_constraint = grammar_constraint35        self.batch_stacks = None36 37    def filter_logits(self, logits, device):38        # resolve each stack to a tensor of True/False for each token39        # indicating acceptance40        # acceptance = self.grammar_acceptor.filter_vocab(self.stacks, device)41        acceptance = self.grammar_constraint.batch_filter_vocab(self.batch_stacks, device)42        # logger.debug(acceptance)43        # Logits to -inf where False44        logits[~acceptance] = -math.inf45 46    # TODO: batching47    def process_logits(self, input_ids, scores, parse_start_index=None):48        """49        :param input_ids:50        :param scores:51        :param parse_start_index: default None, which means generate from scratch. Set to 0 to parse all input_ids52        :return:53        """54        # we dynamically create stacks at the first call, so that we know the batch size and beam size55        if self.batch_stacks is None:56            self.batch_stacks = [self.grammar_constraint.init_stacks() for _ in range(len(input_ids))]57 58        # if self.last_size is not set (which would be the case when processing the first token).59        # In this case, do nothing.60        if self.last_size is None:61            prefix_to_parse = [62                single_input_ids[parse_start_index:] if parse_start_index is not None else []63                for single_input_ids in input_ids64            ]65            # self.grammar_acceptor.accept_token_ids(prefix_to_parse, self.stacks)66            self.batch_stacks = [67                self.grammar_constraint.accept_token_ids(prefix, stack)68                for prefix, stack in zip(prefix_to_parse, self.batch_stacks)69            ]70        #  if the length of the current input IDs (input_ids[0]) is exactly one more than self.last_size.71        #  This is expected in a scenario where inputs are processed incrementally, one token at a time.72        elif len(input_ids[0]) == self.last_size + 1:73            # self.stacks = self.grammar_acceptor.accept_token_id(input_ids[0][-1], self.stacks)74            self.batch_stacks = [75                self.grammar_constraint.accept_token_id(single_input_ids[-1], stack)76                for single_input_ids, stack in zip(input_ids, self.batch_stacks)77            ]78        #  ensure that the input size is consistent with the expected incremental processing79        #  (i.e., one token at a time).80        else:81            # here we check if the input_ids are one token longer than the last time we processed82            # but we don't check if input_ids are actually valid.83            # Imagine a scenario where we generate 10 tokens, then we replace the 10 generated tokens with 10 new tokens.84            # In this case, the input_ids will be consistent with the last_size, but the input_ids are not valid.85            # However, should we really check if the input_ids are valid here?86            # If we do, then we need to reparse the whole input_ids at each call, which is not efficient.87            # Maybe we should just trust the user to provide valid input_ids?88            # The conclusion is that, we assume the input_ids are valid, and our generation will be correct.89            # If the input_ids are not valid, then the generation result will be wrong and we don't take responsibility for that.90            raise RuntimeError(91                "Input ID's length is inconsistent with the current state of "92                "the GrammarConstrainedLogitsProcessor. If you want to process "93                "another input sequence, please instantiate a new "94                "GrammarConstrainedLogitsProcessor."95            )96 97        self.filter_logits(scores, scores.device)98 99        self.last_size = len(input_ids[0])100        return scores101 102    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)103    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:104        return self.process_logits(input_ids, scores)105