codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license class tflogitsprocessor tf method for processing logits raise notimplementederror fself class is an abstract class only classes inheriting this class can be called class tflogitswarper tf method for warping logits raise notimplementederror fself class is an abstract class only classes inheriting this class can be called class tflogitsprocessorlistlist addstartdocstringstflogitsprocessorinputsdocstring def callself inputids tf tensor scores tf tensor curlen int kwargs tf tensor for processor in self functionargs inspect signatureprocessor call parameters if lenfunctionargs 3 if not allarg in kwargs for arg in listfunctionargs keys2 raise valueerror fmake sure that all the required parameters listfunctionargs keys for fprocessor class are passed to the logits processor scores processorinputids scores curlen kwargs else scores processorinputids scores curlen return scores class tftemperaturelogitswarpertflogitswarper r tflogitswarper for temperature exponential scaling output probability distribution args temperature float the value used to module the logits distribution def initself topk int filtervalue float floatinf mintokenstokeep int 1 if not isinstancetopk int or topk 0 raise valueerrorftopk has to be a strictly positive integer but is topk self topk maxtopk mintokenstokeep self filtervalue filtervalue def callself inputids tf tensor scores tf tensor curlen int tf tensor topk minself topk scores shape1 safety check boolean mask containing all tokens with a probability less than the last token of the topk indicestoremove scores tf math topkscores ktopk0 1 nextscores tf whereindicestoremove self filtervalue scores return nextscores class tftopplogitswarpertflogitswarper def initself topp float filtervalue float floatinf mintokenstokeep int 1 if not isinstancetopp float or topp 0 or topp 1 0 raise valueerrorftopp has to be a float 0 and 1 but is topp if not isinstancemintokenstokeep int or mintokenstokeep 1 raise valueerrorfmintokenstokeep has to be a positive integer but is mintokenstokeep self topp topp self filtervalue filtervalue self mintokenstokeep mintokenstokeep def callself inputids tf tensor scores tf tensor curlen int tf tensor topkscores topkindices tf math topkscores scores shape1 maskscores tf fillscores shape self filtervalue cumulativeprobs tf math cumsumstablesoftmaxtopkscores axis1 axis1 scoremask cumulativeprobs self topp also include the token that is higher than topp the first false shift and insert a true on the left scoremask tf concattf onesscoremask shape0 1 dtypetf bool scoremask 1 axis1 ensure min tokens to keep scoremask tf concat tf onesscoremask shape0 self mintokenstokeep dtypetf bool scoremask self mintokenstokeep axis1 mask the values that do not fit the criteria topknextscores tf wherescoremask topkscores maskscores undo the topk sorting converts the 2d matrix of perrow original indices of shape batchsize vocabsize to a 3d tensor of shape batchsize vocabsize 2 containing the original score coordinate from which we can scatter i e scatterindicesrow col is a tensor containing row topkindicesrow col scatterrows tf tiletf expanddimstf rangetopkindices shape0 axis1 1 topkindices shape1 scatterindices tf stackscatterrows topkindices axis1 nextscores tf scatterndscatterindices topknextscores shapetopknextscores shape return nextscores class tfminlengthlogitsprocessortflogitsprocessor r tflogitsprocessor enforcing a minlength by setting eos probability to 0 args minlength int the minimum length below which the score of eostokenid is set to floatinf eostokenid int the id of the endofsequence token applies eos token masking if the first argument is true def initself penalty float if not isinstancepenalty float or not penalty 0 raise valueerrorfpenalty has to be a strictly positive float but is penalty self penalty penalty def createscorepenaltiesself inputids tf tensor logits tf tensor tf tensor we want to populate the penalties in the positions of inputids since xla can t handle shapes unknown before runtime tf unique can t be used therefore we may have redundant updates when a given row has the same token multiple times gathers the penalties to apply logitpenalties tf gatherlogits inputids axis1 batchdims1 logitpenalties tf wherelogitpenalties 0 1 self penalty logitpenalties logitpenalties tf wherelogitpenalties 0 self penalty logitpenalties scatters the penalties tokenpenalties tf oneslogits shape batchsize inputids shape0 seqlen tf shapeinputids1 the sequence length has dynamic size hence the dynamic shape indexableprevinputids tf concat tf expanddimstf repeattf rangebatchsize seqlen axis1 tf expanddimstf reshapeinputids 1 axis1 axis1 tokenpenalties tf tensorscatterndupdate tokenpenalties indicesindexableprevinputids updatestf reshapelogitpenalties 1 return tokenpenalties def callself inputids tf tensor scores tf tensor curlen int tf tensor scorepenalties self createscorepenaltiesinputids curlen scores scores tf math multiplyscores scorepenalties return scores class tfnobadwordslogitsprocessortflogitsprocessor def initself badwordsids listlistint eostokenid int if not isinstancebadwordsids list or lenbadwordsids 0 raise valueerrorfbadwordsids has to be a nonempty list but is badwordsids if anynot isinstancebadwordids list for badwordids in badwordsids raise valueerrorfbadwordsids has to be a list of lists but is badwordsids if any anynot isinstancetokenid int np integer or tokenid 0 for tokenid in badwordids for badwordids in badwordsids raise valueerror feach list in badwordsids has to be a list of positive integers but is badwordsids stores the information about bad words in three tensors 1 a rectangular tensor with the forbidden sequences padded with 1 for full data comparisons self badwordseqsids tf ragged constantbadwordsids totensordefaultvalue1 2 a tensor with the unpadded length of each forbidden sequence for quick length comparisons badwordseqslen lenbadwords for badwords in badwordsids if anywordlen 0 for wordlen in badwordseqslen raise valueerrorfbanned words token sequences badwordsids cannot have an empty list self badwordseqslen tf converttotensorbadwordseqslen dtypetf int32 3 a tensor containing the last token for each sequence for easy access to the tokens that may be banned self seqforbiddentokens tf converttotensorbadwords1 for badwords in badwordsids def calcrowbannedbadtokensself rowinputids tf tensor tf tensor def tokensmatchbadwordseqnumber def lenone if the bad sequence only has one token always mask it return tf cond tf math equalself badwordseqslenbadwordseqnumber 1 lambda tf ones dtypetf bool lengreaterthancurlen def lengreaterthancurlen otherwise if the bad sequence is longer than the current length they can t ever match return tf cond tf math greaterself badwordseqslenbadwordseqnumber tf shaperowinputids0 lambda tf zeros dtypetf bool matchfound def matchfound finaly runs the actual comparison can only be called if the previous comparisons do not yield an answer otherwise we get indexing exceptions comparelen self badwordseqslenbadwordseqnumber 1 return tf cond tf math reduceall tf math equal rowinputidscomparelen self badwordseqsidsbadwordseqnumber comparelen lambda tf ones dtypetf bool lambda tf zeros dtypetf bool match lenone return match compares the current row against all bad word sequences obtaining a mask with the matches matchmask tf mapfntokensmatch tf rangeself badwordseqsids shape0 fnoutputsignaturetf bool rowbannedtokens self seqforbiddentokensmatchmask return rowbannedtokens def callself inputids tf tensor scores tf tensor curlen int tf tensor we want to mask some banned tokens at a score level since the banned tokens depend on the previous inputids they may have a different length for each row and they may even be empty for some rows to remain simple and xlacompatible we work on a perrow fashion todo joao this function might trigger xla retracing as curlen increases fix it if it becomes a frequent choke point make curlen a tensor def getrowupdatedscorerowinputs tupletf tensor tf tensor rowinputids rowscore rowinputs bannedtokens self calcrowbannedbadtokensrowinputids curlen bannedtokensmask tf scatternd indicestf expanddimsbannedtokens axis1 updatestf oneslikebannedtokens dtypetf bool shaperowscore shape rowscore tf wherebannedtokensmask floatinf rowscore return rowscore scores tf mapfngetrowupdatedscore inputids scores fnoutputsignaturetf float32 return scores class tfnorepeatngramlogitsprocessortflogitsprocessor r tflogitsprocessor that enforces no repetition of ngrams see fairseqhttps github compytorchfairseqbloba07cb6f40480928c9e0548b737aadd36ee66ac76fairseqsequencegenerator pyl345 args ngramsize int all ngrams of size ngramsize can only occur once copied from fairseq for norepeatngram in beamsearch return no banned tokens if we haven t generated ngramsize tokens yet before decoding the next token prevent decoding of ngrams that have already appeared todo joao enable xla on this logits processor see discussion and attempts in https github comhuggingfacetransformerspull16974 create bannedtokens boolean mask def initself bostokenid int if bostokenid 0 raise valueerrorfthe forced bos token id must be a nonnegative integer got bostokenid self bostokenid bostokenid def callself inputids tf tensor scores tf tensor curlen int tf tensor if curlen 1 batchsize numtokens scores shape sets the score to 0 in the bostokenid column scores tf zerosbatchsize 1 sets the score to inf everywhere else if self bostokenid 0 scores tf concattf broadcasttofloatinf batchsize self bostokenid scores axis1 if self bostokenid numtokens 1 scores tf concat scores tf broadcasttofloatinf batchsize numtokens 1 self bostokenid axis1 return scores class tfforcedeostokenlogitsprocessortflogitsprocessor r tflogitsprocessor that enforces the specified token as the last generated token when maxlength is reached args maxlength int the maximum length of the sequence to be generated eostokenid int the id of the token to force as the last generated token when maxlength is reached sets the score to 0 in the eostokenid column sets the score to inf everywhere else def initself beginsuppresstokens beginindex self beginsuppresstokens listbeginsuppresstokens self beginindex beginindex def callself inputids tf tensor scores tf tensor curlen int tf tensor scores tf cond tf equalcurlen self beginindex lambda tf tensorscatterndupdate scores indicesi token for i in rangescores shape0 for token in self beginsuppresstokens updatesfloatinf for in rangescores shape0 lenself beginsuppresstokens lambda scores return scores class tfsuppresstokenslogitsprocessortflogitsprocessor rthis processor can be used to suppress a list of tokens the processor will set their log probs to inf so that they are not sampled def initself suppresstokens self suppresstokens listsuppresstokens def callself inputids tf tensor scores tf tensor curlen int tf tensor scores tf tensorscatterndupdate scores indicesi token for i in rangescores shape0 for token in self suppresstokens updatesfloatinf for in rangescores shape0 lenself suppresstokens return scores class tfforcetokenslogitsprocessortflogitsprocessor rthis processor takes a list of pairs of integers which indicates a mapping from generation indices to token indices that will be forced before sampling the processor will set their log probs to 0 and all other tokens to inf so that they are sampled at their corresponding index def initself forcetokenmap listlistint forcetokenmap dictforcetokenmap converts the dictionary of format index token containing the tokens to be forced to an array where the index of the array corresponds to the index of the token to be forced for xla compatibility indexes without forced tokens will have an negative value forcetokenarray np onesmaxforcetokenmap keys 1 dtypenp int32 1 for index token in forcetokenmap items if token is not none forcetokenarrayindex token self forcetokenarray tf converttotensorforcetokenarray dtypetf int32 def callself inputids tf tensor scores tf tensor curlen int tf tensor def forcetokengenerationidx batchsize scores shape0 currenttoken self forcetokenarraygenerationidx newscores tf oneslikescores dtypescores dtype floatinf indices tf stacktf rangebatchsize tf tilecurrenttoken batchsize axis1 updates tf zerosbatchsize dtypescores dtype newscores tf tensorscatterndupdatenewscores indices updates return newscores scores tf cond tf greaterequalcurlen tf shapeself forcetokenarray0 if the current length is geq than the length of forcetokenarray the processor does nothing lambda tf identityscores otherwise it may force a certain token lambda tf cond tf greaterequalself forcetokenarraycurlen 0 only valid positive tokens are forced lambda forcetokencurlen otherwise the processor does nothing lambda scores return scores coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license args input_ids tf tensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary indices can be obtained using pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids scores tf tensor of shape batch_size config vocab_size prediction scores of a language modeling head these can be logits for each vocabulary when not using beam search or log softmax for each vocabulary token when using beam search cur_len int the current length of valid input sequence tokens in the tf implementation the input_ids sequence length is the maximum length generate can produce and we need to know which of its tokens are valid kwargs dict str any optional additional logits processor specific kwargs return tf tensor of shape batch_size config vocab_size the processed prediction scores abstract base class for all logit processors that can be applied during generation tf method for processing logits abstract base class for all logit warpers that can be applied during generation with multinomial sampling tf method for warping logits this class can be used to create a list of tflogitsprocessor to subsequently process a scores input tensor this class inherits from list and adds a specific __call__ method to apply each tflogitsprocessor to the inputs tflogitswarper for temperature exponential scaling output probability distribution args temperature float the value used to module the logits distribution tflogitswarper that performs top k i e restricting to the k highest probability elements args top_k int the number of highest probability vocabulary tokens to keep for top k filtering filter_value float optional defaults to inf all filtered values will be set to this float value min_tokens_to_keep int optional defaults to 1 minimum number of tokens that cannot be filtered safety check boolean mask containing all tokens with a probability less than the last token of the top k tflogitswarper that performs top p i e restricting to top tokens summing to prob_cut_off args top_p float if set to 1 only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation filter_value float optional defaults to inf all filtered values will be set to this float value min_tokens_to_keep int optional defaults to 1 minimum number of tokens that cannot be filtered also include the token that is higher than top_p the first false shift and insert a true on the left ensure min tokens to keep mask the values that do not fit the criteria undo the topk sorting converts the 2d matrix of per row original indices of shape batch_size vocab_size to a 3d tensor of shape batch_size vocab_size 2 containing the original score coordinate from which we can scatter i e scatter_indices row col is a tensor containing row topk_indices row col tflogitsprocessor enforcing a min length by setting eos probability to 0 args min_length int the minimum length below which the score of eos_token_id is set to float inf eos_token_id int the id of the end of sequence token applies eos token masking if the first argument is true tflogitsprocessor enforcing an exponential penalty on repeated sequences args repetition_penalty float the parameter for repetition penalty 1 0 means no penalty see this paper https arxiv org pdf 1909 05858 pdf for more details we want to populate the penalties in the positions of input_ids since xla can t handle shapes unknown before runtime tf unique can t be used therefore we may have redundant updates when a given row has the same token multiple times gathers the penalties to apply scatters the penalties the sequence length has dynamic size hence the dynamic shape tflogitsprocessor that enforces that specified sequences will never be sampled args bad_words_ids list list int list of list of token ids that are not allowed to be generated in order to get the tokens of the words that should not appear in the generated text make sure to set add_prefix_space true when initializing the tokenizer and use tokenizer bad_words add_special_tokens false input_ids the add_prefix_space argument is only supported for some slow tokenizers as fast tokenizers prefixing behaviours come from pre tokenizers read more here https huggingface co docs tokenizers api pre tokenizers eos_token_id int the id of the end of sequence token stores the information about bad words in three tensors 1 a rectangular tensor with the forbidden sequences padded with 1 for full data comparisons 2 a tensor with the unpadded length of each forbidden sequence for quick length comparisons 3 a tensor containing the last token for each sequence for easy access to the tokens that may be banned if the bad sequence only has one token always mask it otherwise if the bad sequence is longer than the current length they can t ever match finaly runs the actual comparison can only be called if the previous comparisons do not yield an answer otherwise we get indexing exceptions compares the current row against all bad word sequences obtaining a mask with the matches we want to mask some banned tokens at a score level since the banned tokens depend on the previous input_ids they may have a different length for each row and they may even be empty for some rows to remain simple and xla compatible we work on a per row fashion todo joao this function might trigger xla retracing as cur_len increases fix it if it becomes a frequent choke point make cur_len a tensor tflogitsprocessor that enforces no repetition of n grams see fairseq https github com pytorch fairseq blob a07cb6f40480928c9e0548b737aadd36ee66ac76 fairseq sequence_generator py l345 args ngram_size int all ngrams of size ngram_size can only occur once copied from fairseq for no_repeat_ngram in beam_search return no banned tokens if we haven t generated ngram_size tokens yet before decoding the next token prevent decoding of ngrams that have already appeared todo joao enable xla on this logits processor see discussion and attempts in https github com huggingface transformers pull 16974 create banned_tokens boolean mask tflogitsprocessor that enforces the specified token as the first generated token args bos_token_id int the id of the token to force as the first generated token sets the score to 0 in the bos_token_id column sets the score to inf everywhere else tflogitsprocessor that enforces the specified token as the last generated token when max_length is reached args max_length int the maximum length of the sequence to be generated eos_token_id int the id of the token to force as the last generated token when max_length is reached sets the score to 0 in the eos_token_id column sets the score to inf everywhere else tfsuppresstokensatbeginlogitsprocessor suppresses a list of tokens as soon as the generate function starts generating using begin_index tokens this should ensure that the tokens defined by begin_suppress_tokens at not sampled at the begining of the generation this processor can be used to suppress a list of tokens the processor will set their log probs to inf so that they are not sampled this processor takes a list of pairs of integers which indicates a mapping from generation indices to token indices that will be forced before sampling the processor will set their log probs to 0 and all other tokens to inf so that they are sampled at their corresponding index converts the dictionary of format index token containing the tokens to be forced to an array where the index of the array corresponds to the index of the token to be forced for xla compatibility indexes without forced tokens will have an negative value if the current length is geq than the length of force_token_array the processor does nothing otherwise it may force a certain token only valid positive tokens are forced otherwise the processor does nothing
import inspect
from typing import List, Tuple

import numpy as np
import tensorflow as tf

from ..tf_utils import stable_softmax
from ..utils import add_start_docstrings
from ..utils.logging import get_logger


logger = get_logger(__name__)


TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING = r


class TFLogitsProcessor:
    

    @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )


class TFLogitsWarper:
    

    @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )


class TFLogitsProcessorList(list):
    

    @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int, **kwargs) -> tf.Tensor:
        for processor in self:
            function_args = inspect.signature(processor.__call__).parameters
            if len(function_args) > 3:
                if not all(arg in kwargs for arg in list(function_args.keys())[2:]):
                    raise ValueError(
                        f"Make sure that all the required parameters: {list(function_args.keys())} for "
                        f"{processor.__class__} are passed to the logits processor."
                    )
                scores = processor(input_ids, scores, cur_len, **kwargs)
            else:
                scores = processor(input_ids, scores, cur_len)
        return scores


class TFTemperatureLogitsWarper(TFLogitsWarper):
    r

    def __init__(self, temperature: float):
        if not isinstance(temperature, float) or not (temperature > 0):
            raise ValueError(f"`temperature` has to be a strictly positive float, but is {temperature}")

        self.temperature = temperature

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        scores = scores / self.temperature
        return scores


class TFTopKLogitsWarper(TFLogitsWarper):
    r

    def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        if not isinstance(top_k, int) or top_k <= 0:
            raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}")

        self.top_k = max(top_k, min_tokens_to_keep)
        self.filter_value = filter_value

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        top_k = min(self.top_k, scores.shape[-1])  
        
        indices_to_remove = scores < tf.math.top_k(scores, k=top_k)[0][..., -1:]
        next_scores = tf.where(indices_to_remove, self.filter_value, scores)
        return next_scores


class TFTopPLogitsWarper(TFLogitsWarper):
    

    def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        if not isinstance(top_p, float) or (top_p < 0 or top_p > 1.0):
            raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")
        if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
            raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")

        self.top_p = top_p
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        topk_scores, topk_indices = tf.math.top_k(scores, scores.shape[-1])

        mask_scores = tf.fill(scores.shape, self.filter_value)
        cumulative_probs = tf.math.cumsum(stable_softmax(topk_scores, axis=-1), axis=-1)
        score_mask = cumulative_probs < self.top_p

        
        score_mask = tf.concat((tf.ones([score_mask.shape[0], 1], dtype=tf.bool), score_mask[:, :-1]), axis=-1)

        
        score_mask = tf.concat(
            (
                tf.ones([score_mask.shape[0], self.min_tokens_to_keep], dtype=tf.bool),
                score_mask[:, self.min_tokens_to_keep :],
            ),
            axis=-1,
        )

        
        topk_next_scores = tf.where(score_mask, topk_scores, mask_scores)

        
        
        
        scatter_rows = tf.tile(tf.expand_dims(tf.range(topk_indices.shape[0]), axis=-1), [1, topk_indices.shape[-1]])
        scatter_indices = tf.stack((scatter_rows, topk_indices), axis=-1)
        next_scores = tf.scatter_nd(scatter_indices, topk_next_scores, shape=topk_next_scores.shape)

        return next_scores


class TFMinLengthLogitsProcessor(TFLogitsProcessor):
    r

    def __init__(self, min_length: int, eos_token_id: int):
        if not isinstance(min_length, int) or min_length < 0:
            raise ValueError(f"`min_length` has to be a positive integer, but is {min_length}")

        if not isinstance(eos_token_id, int) or eos_token_id < 0:
            raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}")

        self.min_length = min_length
        self.eos_token_id = eos_token_id

    def _apply_eos_token_mask(self, scores: tf.Tensor) -> tf.Tensor:
        eos_token_id_mask = tf.range(scores.shape[-1]) == self.eos_token_id
        scores = tf.where(eos_token_id_mask, float("-inf"), scores)
        return scores

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        
        scores = tf.cond(
            tf.less(cur_len, self.min_length),
            lambda: self._apply_eos_token_mask(scores),
            lambda: tf.identity(scores),
        )
        return scores


class TFRepetitionPenaltyLogitsProcessor(TFLogitsProcessor):
    r

    def __init__(self, penalty: float):
        if not isinstance(penalty, float) or not (penalty > 0):
            raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}")

        self.penalty = penalty

    def _create_score_penalties(self, input_ids: tf.Tensor, logits: tf.Tensor) -> tf.Tensor:
        
        
        

        
        logit_penalties = tf.gather(logits, input_ids, axis=1, batch_dims=1)
        logit_penalties = tf.where(logit_penalties > 0, 1 / self.penalty, logit_penalties)
        logit_penalties = tf.where(logit_penalties < 0, self.penalty, logit_penalties)

        
        token_penalties = tf.ones(logits.shape)
        batch_size = input_ids.shape[0]
        seq_len = tf.shape(input_ids)[1]  
        indexable_prev_input_ids = tf.concat(
            (
                tf.expand_dims(tf.repeat(tf.range(batch_size), seq_len), axis=-1),
                tf.expand_dims(tf.reshape(input_ids, [-1]), axis=-1),
            ),
            axis=1,
        )
        token_penalties = tf.tensor_scatter_nd_update(
            token_penalties, indices=indexable_prev_input_ids, updates=tf.reshape(logit_penalties, [-1])
        )
        return token_penalties

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        score_penalties = self._create_score_penalties(input_ids[:, :cur_len], scores)

        scores = tf.math.multiply(scores, score_penalties)

        return scores


class TFNoBadWordsLogitsProcessor(TFLogitsProcessor):
    

    def __init__(self, bad_words_ids: List[List[int]], eos_token_id: int):
        if not isinstance(bad_words_ids, List) or len(bad_words_ids) == 0:
            raise ValueError(f"`bad_words_ids` has to be a non-empty list, but is {bad_words_ids}.")
        if any(not isinstance(bad_word_ids, list) for bad_word_ids in bad_words_ids):
            raise ValueError(f"`bad_words_ids` has to be a list of lists, but is {bad_words_ids}.")
        if any(
            any((not isinstance(token_id, (int, np.integer)) or token_id < 0) for token_id in bad_word_ids)
            for bad_word_ids in bad_words_ids
        ):
            raise ValueError(
                f"Each list in `bad_words_ids` has to be a list of positive integers, but is {bad_words_ids}."
            )

        
        
        self.bad_word_seqs_ids = tf.ragged.constant(bad_words_ids).to_tensor(default_value=-1)
        
        bad_word_seqs_len = [len(bad_words) for bad_words in bad_words_ids]
        if any(word_len == 0 for word_len in bad_word_seqs_len):
            raise ValueError(f"Banned words token sequences {bad_words_ids} cannot have an empty list")
        self.bad_word_seqs_len = tf.convert_to_tensor(bad_word_seqs_len, dtype=tf.int32)
        
        self.seq_forbidden_tokens = tf.convert_to_tensor([bad_words[-1] for bad_words in bad_words_ids])

    def _calc_row_banned_bad_tokens(self, row_input_ids: tf.Tensor) -> tf.Tensor:
        def _tokens_match(bad_word_seq_number):
            def _len_one():
                
                return tf.cond(
                    tf.math.equal(self.bad_word_seqs_len[bad_word_seq_number], 1),
                    lambda: tf.ones((), dtype=tf.bool),
                    _len_greater_than_cur_len,
                )

            def _len_greater_than_cur_len():
                
                return tf.cond(
                    tf.math.greater(self.bad_word_seqs_len[bad_word_seq_number], tf.shape(row_input_ids)[0]),
                    lambda: tf.zeros((), dtype=tf.bool),
                    _match_found,
                )

            def _match_found():
                
                
                compare_len = self.bad_word_seqs_len[bad_word_seq_number] - 1
                return tf.cond(
                    tf.math.reduce_all(
                        tf.math.equal(
                            row_input_ids[-compare_len:], self.bad_word_seqs_ids[bad_word_seq_number, :compare_len]
                        )
                    ),
                    lambda: tf.ones((), dtype=tf.bool),
                    lambda: tf.zeros((), dtype=tf.bool),
                )

            match = _len_one()
            return match

        
        match_mask = tf.map_fn(_tokens_match, tf.range(self.bad_word_seqs_ids.shape[0]), fn_output_signature=tf.bool)
        row_banned_tokens = self.seq_forbidden_tokens[match_mask]
        return row_banned_tokens

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        
        
        
        
        
        def _get_row_updated_score(row_inputs: Tuple[tf.Tensor]) -> tf.Tensor:
            row_input_ids, row_score = row_inputs
            banned_tokens = self._calc_row_banned_bad_tokens(row_input_ids[:cur_len])
            banned_tokens_mask = tf.scatter_nd(
                indices=tf.expand_dims(banned_tokens, axis=-1),
                updates=tf.ones_like(banned_tokens, dtype=tf.bool),
                shape=row_score.shape,
            )
            row_score = tf.where(banned_tokens_mask, -float("inf"), row_score)
            return row_score

        scores = tf.map_fn(_get_row_updated_score, (input_ids, scores), fn_output_signature=tf.float32)
        return scores


class TFNoRepeatNGramLogitsProcessor(TFLogitsProcessor):
    r

    def __init__(self, ngram_size: int):
        if not isinstance(ngram_size, int) or ngram_size <= 0:
            raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}")
        self.ngram_size = ngram_size

    def calc_banned_ngram_tokens(self, input_ids, num_hypos, cur_len):
        
        if cur_len + 1 < self.ngram_size:
            
            return [[] for _ in range(num_hypos)]
        generated_ngrams = [{} for _ in range(num_hypos)]
        prev_input_ids = input_ids[:, :cur_len]
        for idx in range(num_hypos):
            gen_tokens = prev_input_ids[idx].numpy().tolist()
            generated_ngram = generated_ngrams[idx]
            for ngram in zip(*[gen_tokens[i:] for i in range(self.ngram_size)]):
                prev_ngram_tuple = tuple(ngram[:-1])
                generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]]

        def _get_generated_ngrams(hypo_idx):
            
            start_idx = cur_len + 1 - self.ngram_size
            ngram_idx = tuple(prev_input_ids[hypo_idx, start_idx:cur_len].numpy().tolist())
            return generated_ngrams[hypo_idx].get(ngram_idx, [])

        banned_tokens = [_get_generated_ngrams(hypo_idx) for hypo_idx in range(num_hypos)]

        return banned_tokens

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        
        
        if not tf.executing_eagerly():
            raise NotImplementedError("TFNoRepeatNGramLogitsProcessor is only implemented for eager execution.")

        batch_size, vocab_size = scores.shape
        banned_tokens = self.calc_banned_ngram_tokens(input_ids, batch_size, cur_len)

        
        banned_tokens_indices_mask = []
        for banned_tokens_slice in banned_tokens:
            banned_tokens_indices_mask.append(
                [True if token in banned_tokens_slice else False for token in range(vocab_size)]
            )

        scores = tf.where(tf.convert_to_tensor(banned_tokens_indices_mask, dtype=tf.bool), -float("inf"), scores)

        return scores


class TFForcedBOSTokenLogitsProcessor(TFLogitsProcessor):
    r

    def __init__(self, bos_token_id: int):
        if bos_token_id < 0:
            raise ValueError(f"The forced bos token id  must be a non-negative integer, got {bos_token_id}")
        self.bos_token_id = bos_token_id

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        if cur_len == 1:
            batch_size, num_tokens = scores.shape
            
            scores = tf.zeros((batch_size, 1))
            
            if self.bos_token_id > 0:
                scores = tf.concat((tf.broadcast_to(-float("inf"), (batch_size, self.bos_token_id)), scores), axis=-1)
            if self.bos_token_id < (num_tokens - 1):
                scores = tf.concat(
                    (scores, tf.broadcast_to(-float("inf"), (batch_size, (num_tokens - 1) - self.bos_token_id))),
                    axis=-1,
                )
        return scores


class TFForcedEOSTokenLogitsProcessor(TFLogitsProcessor):
    r

    def __init__(self, max_length: int, eos_token_id: int):
        self.max_length = max_length
        if eos_token_id < 0:
            raise ValueError(f"The forced eos token id must be a non-negative integer, got {eos_token_id}")
        self.eos_token_id = eos_token_id

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        if cur_len == self.max_length - 1:
            batch_size, num_tokens = scores.shape
            
            scores = tf.zeros((batch_size, 1))
            
            if self.eos_token_id > 0:
                scores = tf.concat((tf.broadcast_to(-float("inf"), (batch_size, self.eos_token_id)), scores), axis=-1)
            if self.eos_token_id < (num_tokens - 1):
                scores = tf.concat(
                    (scores, tf.broadcast_to(-float("inf"), (batch_size, (num_tokens - 1) - self.eos_token_id))),
                    axis=-1,
                )
        return scores


class TFSuppressTokensAtBeginLogitsProcessor(TFLogitsProcessor):
    r

    def __init__(self, begin_suppress_tokens, begin_index):
        self.begin_suppress_tokens = list(begin_suppress_tokens)
        self.begin_index = begin_index

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        scores = tf.cond(
            tf.equal(cur_len, self.begin_index),
            lambda: tf.tensor_scatter_nd_update(
                scores,
                indices=[[i, token] for i in range(scores.shape[0]) for token in self.begin_suppress_tokens],
                updates=[-float("inf") for _ in range(scores.shape[0] * len(self.begin_suppress_tokens))],
            ),
            lambda: scores,
        )
        return scores


class TFSuppressTokensLogitsProcessor(TFLogitsProcessor):
    r

    def __init__(self, suppress_tokens):
        self.suppress_tokens = list(suppress_tokens)

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        scores = tf.tensor_scatter_nd_update(
            scores,
            indices=[[i, token] for i in range(scores.shape[0]) for token in self.suppress_tokens],
            updates=[-float("inf") for _ in range(scores.shape[0] * len(self.suppress_tokens))],
        )
        return scores


class TFForceTokensLogitsProcessor(TFLogitsProcessor):
    r

    def __init__(self, force_token_map: List[List[int]]):
        force_token_map = dict(force_token_map)
        
        
        
        force_token_array = np.ones((max(force_token_map.keys()) + 1), dtype=np.int32) * -1
        for index, token in force_token_map.items():
            if token is not None:
                force_token_array[index] = token
        self.force_token_array = tf.convert_to_tensor(force_token_array, dtype=tf.int32)

    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
        def _force_token(generation_idx):
            batch_size = scores.shape[0]
            current_token = self.force_token_array[generation_idx]

            new_scores = tf.ones_like(scores, dtype=scores.dtype) * -float("inf")
            indices = tf.stack((tf.range(batch_size), tf.tile([current_token], [batch_size])), axis=1)
            updates = tf.zeros((batch_size,), dtype=scores.dtype)
            new_scores = tf.tensor_scatter_nd_update(new_scores, indices, updates)
            return new_scores

        scores = tf.cond(
            tf.greater_equal(cur_len, tf.shape(self.force_token_array)[0]),
            
            lambda: tf.identity(scores),
            
            lambda: tf.cond(
                tf.greater_equal(self.force_token_array[cur_len], 0),
                
                lambda: _force_token(cur_len),
                
                lambda: scores,
            ),
        )
        return scores