codingutf8 2021 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 flaxlogitsprocessor flax method for processing logits raise notimplementederror fself class is an abstract class only classes inheriting this class can be called class flaxlogitswarper flax method for warping logits raise notimplementederror fself class is an abstract class only classes inheriting this class can be called class flaxlogitsprocessorlistlist addstartdocstringslogitsprocessorinputsdocstring def callself inputids jnp ndarray scores jnp ndarray curlen int kwargs jnp ndarray 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 flaxtemperaturelogitswarperflaxlogitswarper r flaxlogitswarper for temperature exponential scaling output probability distribution args temperature float the value used to module the logits distribution flaxlogitswarper that performs topp i e restricting to top tokens summing to probcutoff probcutoff args topp float if set to 1 only the smallest set of most probable tokens with probabilities that add up to topp or higher are kept for generation filtervalue float optional defaults to inf all filtered values will be set to this float value mintokenstokeep int optional defaults to 1 minimum number of tokens that cannot be filtered include the token that is higher than topp as well min tokens to keep 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 jnp ndarray scores jnp ndarray curlen int jnp ndarray batchsize vocabsize scores shape nextscoresflat jnp fullbatchsize vocabsize self filtervalue topk minself topk scores shape1 safety check topkscores topkindices lax topkscores topk shift jnp broadcasttojnp arangebatchsize vocabsize none batchsize topk flatten topkscoresflat topkscores flatten topkindicesflat topkindices flatten shift nextscoresflat nextscoresflat attopkindicesflat settopkscoresflat nextscores nextscoresflat reshapebatchsize vocabsize return nextscores class flaxforcedbostokenlogitsprocessorflaxlogitsprocessor r flaxlogitsprocessor that enforces the specified token as the first generated token args bostokenid int the id of the token to force as the first generated token def initself maxlength int eostokenid int self maxlength maxlength self eostokenid eostokenid def callself inputids jnp ndarray scores jnp ndarray curlen int jnp ndarray newscores jnp fullscores shape floatinf applypenalty 1 jnp boolcurlen self maxlength 1 scores jnp whereapplypenalty newscores at self eostokenid set0 scores return scores class flaxminlengthlogitsprocessorflaxlogitsprocessor r flaxlogitsprocessor 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 create boolean flag to decide if min length penalty should be applied def initself beginsuppresstokens beginindex self beginsuppresstokens listbeginsuppresstokens self beginindex beginindex def callself inputids scores curlen int applypenalty 1 jnp boolcurlen self beginindex scores jnp whereapplypenalty scores at self beginsuppresstokens setfloatinf scores return scores class flaxsuppresstokenslogitsprocessorflaxlogitsprocessor r flaxlogitsprocessor suppressing a list of tokens at each decoding step the processor will set their log probs to be inf so they are not sampled args suppresstokens list tokens to not sample def initself forcetokenmap 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 a negative value forcetokenarray jnp onesmaxforcetokenmap keys 1 dtypejnp int32 1 for index token in forcetokenmap items if token is not none forcetokenarray forcetokenarray atindex settoken self forcetokenarray jnp int32forcetokenarray def callself inputids jnp ndarray scores jnp ndarray curlen int jnp ndarray def forcetokengenerationidx batchsize scores shape0 currenttoken self forcetokenarraygenerationidx newscores jnp oneslikescores dtypescores dtype floatinf updates jnp zerosbatchsize 1 dtypescores dtype newscores lax dynamicupdateslicenewscores updates 0 currenttoken return newscores scores lax cond curlen self forcetokenarray shape0 if the current length is geq than the length of forcetokenarray the processor does nothing lambda scores otherwise it may force a certain token lambda lax cond self forcetokenarraycurlen 0 only valid positive tokens are forced lambda forcetokencurlen otherwise the processor does nothing lambda scores return scores class flaxwhispertimestamplogitsprocessorflaxlogitsprocessor r whisper specific processor this processor can be used to force a list of tokens the processor will set their log probs to inf so that they are sampled at their corresponding index args generateconfig generateconfig the generate config used to generate the output the following parameters are required eostokenid int optional defaults to 50257 the id of the endofsequence token notimestampstokenid int optional defaults to 50363 the id of the notimestamps token maxinitialtimestampindex int optional defaults to 1 used to set the maximum value of the initial timestamp this is used to prevent the model from predicting timestamps that are too far in the future room for language token and task token suppress notimestamps which is handled by withouttimestamps if sum of probability over timestamps is above any other token sample timestamp coding utf 8 2021 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 jnp ndarray 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 jnp ndarray 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 kwargs dict str any optional additional logits processor specific kwargs return jnp ndarray of shape batch_size config vocab_size the processed prediction scores abstract base class for all logit processors that can be applied during generation flax method for processing logits abstract base class for all logit warpers that can be applied during generation with multinomial sampling flax method for warping logits this class can be used to create a list of flaxlogitsprocessor or flaxlogitswarper to subsequently process a scores input tensor this class inherits from list and adds a specific __call__ method to apply each flaxlogitsprocessor or flaxlogitswarper to the inputs flaxlogitswarper for temperature exponential scaling output probability distribution args temperature float the value used to module the logits distribution flaxlogitswarper that performs top p i e restricting to top tokens summing to prob_cut_off 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 include the token that is higher than top_p as well min tokens to keep flaxlogitswarper 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 flaxlogitsprocessor 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 flaxlogitsprocessor 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 flaxlogitsprocessor 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 create boolean flag to decide if min length penalty should be applied flaxlogitsprocessor supressing 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 are not sampled at the begining of the generation args begin_suppress_tokens list int tokens to not sample begin_index int index where the tokens are suppressed flaxlogitsprocessor suppressing a list of tokens at each decoding step the processor will set their log probs to be inf so they are not sampled args suppress_tokens list tokens to not sample flaxlogitsprocessor that 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 args force_token_map list map giving token ids and indices where they will be forced to be sampled 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 a 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 whisper specific processor this processor can be used to force a list of tokens the processor will set their log probs to inf so that they are sampled at their corresponding index args generate_config generateconfig the generate config used to generate the output the following parameters are required eos_token_id int optional defaults to 50257 the id of the end of sequence token no_timestamps_token_id int optional defaults to 50363 the id of the notimestamps token max_initial_timestamp_index int optional defaults to 1 used to set the maximum value of the initial timestamp this is used to prevent the model from predicting timestamps that are too far in the future room for language token and task token suppress notimestamps which is handled by without_timestamps if sum of probability over timestamps is above any other token sample timestamp
import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


logger = get_logger(__name__)


LOGITS_PROCESSOR_INPUTS_DOCSTRING = r


class FlaxLogitsProcessor:
    

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


class FlaxLogitsWarper:
    

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


class FlaxLogitsProcessorList(list):
    

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int, **kwargs) -> jnp.ndarray:
        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 FlaxTemperatureLogitsWarper(FlaxLogitsWarper):
    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: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        scores = scores / self.temperature
        return scores


class FlaxTopPLogitsWarper(FlaxLogitsWarper):
    

    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: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        topk_scores, topk_indices = lax.top_k(scores, scores.shape[-1])

        mask_scores = jnp.full_like(scores, self.filter_value)
        cumulative_probs = jax.nn.softmax(topk_scores, axis=-1).cumsum(axis=-1)
        score_mask = cumulative_probs < self.top_p

        
        score_mask = jnp.roll(score_mask, 1)
        score_mask |= score_mask.at[:, 0].set(True)

        
        score_mask = score_mask.at[:, : self.min_tokens_to_keep].set(True)

        topk_next_scores = jnp.where(score_mask, topk_scores, mask_scores)
        next_scores = jax.lax.sort_key_val(topk_indices, topk_next_scores)[-1]

        return next_scores


class FlaxTopKLogitsWarper(FlaxLogitsWarper):
    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: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        batch_size, vocab_size = scores.shape
        next_scores_flat = jnp.full(batch_size * vocab_size, self.filter_value)

        topk = min(self.top_k, scores.shape[-1])  
        topk_scores, topk_indices = lax.top_k(scores, topk)
        shift = jnp.broadcast_to((jnp.arange(batch_size) * vocab_size)[:, None], (batch_size, topk)).flatten()
        topk_scores_flat = topk_scores.flatten()
        topk_indices_flat = topk_indices.flatten() + shift

        next_scores_flat = next_scores_flat.at[topk_indices_flat].set(topk_scores_flat)
        next_scores = next_scores_flat.reshape(batch_size, vocab_size)
        return next_scores


class FlaxForcedBOSTokenLogitsProcessor(FlaxLogitsProcessor):
    r

    def __init__(self, bos_token_id: int):
        self.bos_token_id = bos_token_id

    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        new_scores = jnp.full(scores.shape, -float("inf"))

        apply_penalty = 1 - jnp.bool_(cur_len - 1)

        scores = jnp.where(apply_penalty, new_scores.at[:, self.bos_token_id].set(0), scores)

        return scores


class FlaxForcedEOSTokenLogitsProcessor(FlaxLogitsProcessor):
    r

    def __init__(self, max_length: int, eos_token_id: int):
        self.max_length = max_length
        self.eos_token_id = eos_token_id

    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        new_scores = jnp.full(scores.shape, -float("inf"))

        apply_penalty = 1 - jnp.bool_(cur_len - self.max_length + 1)

        scores = jnp.where(apply_penalty, new_scores.at[:, self.eos_token_id].set(0), scores)

        return scores


class FlaxMinLengthLogitsProcessor(FlaxLogitsProcessor):
    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 __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        
        apply_penalty = 1 - jnp.clip(cur_len - self.min_length, 0, 1)

        scores = jnp.where(apply_penalty, scores.at[:, self.eos_token_id].set(-float("inf")), scores)

        return scores


class FlaxSuppressTokensAtBeginLogitsProcessor(FlaxLogitsProcessor):
    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, scores, cur_len: int):
        apply_penalty = 1 - jnp.bool_(cur_len - self.begin_index)

        scores = jnp.where(apply_penalty, scores.at[:, self.begin_suppress_tokens].set(-float("inf")), scores)

        return scores


class FlaxSuppressTokensLogitsProcessor(FlaxLogitsProcessor):
    r

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

    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        scores = scores.at[..., self.suppress_tokens].set(-float("inf"))

        return scores


class FlaxForceTokensLogitsProcessor(FlaxLogitsProcessor):
    r

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

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

            new_scores = jnp.ones_like(scores, dtype=scores.dtype) * -float("inf")
            updates = jnp.zeros((batch_size, 1), dtype=scores.dtype)
            new_scores = lax.dynamic_update_slice(new_scores, updates, (0, current_token))
            return new_scores

        scores = lax.cond(
            cur_len >= self.force_token_array.shape[0],
            
            lambda: scores,
            
            lambda: lax.cond(
                self.force_token_array[cur_len] >= 0,
                
                lambda: _force_token(cur_len),
                
                lambda: scores,
            ),
        )
        return scores


class FlaxWhisperTimeStampLogitsProcessor(FlaxLogitsProcessor):
    r

    def __init__(self, generate_config, model_config, decoder_input_length):
        self.eos_token_id = generate_config.eos_token_id
        self.no_timestamps_token_id = generate_config.no_timestamps_token_id
        self.timestamp_begin = generate_config.no_timestamps_token_id + 1

        self.begin_index = decoder_input_length + 1

        if generate_config.is_multilingual:
            
            self.begin_index += 2
        if hasattr(generate_config, "max_initial_timestamp_index"):
            self.max_initial_timestamp_index = generate_config.max_initial_timestamp_index
        else:
            self.max_initial_timestamp_index = model_config.vocab_size
        if self.max_initial_timestamp_index is None:
            self.max_initial_timestamp_index = model_config.vocab_size

    def __call__(self, input_ids, scores, cur_len):
        
        scores = scores.at[:, self.no_timestamps_token_id].set(-float("inf"))

        def handle_pairs(input_ids_k, scores_k):
            last_was_timestamp = jnp.where((cur_len - self.begin_index) >= 1, True, False)
            last_was_timestamp = jnp.where(
                input_ids_k[cur_len - 1] >= self.timestamp_begin,
                True and last_was_timestamp,
                False,
            )

            penultimate_was_timestamp = jnp.where((cur_len - self.begin_index) < 2, True, False)
            penultimate_was_timestamp = jnp.where(
                input_ids_k[cur_len - 2] >= self.timestamp_begin,
                True,
                penultimate_was_timestamp,
            )

            return jnp.where(
                last_was_timestamp,
                jnp.where(
                    penultimate_was_timestamp > 0,
                    scores_k.at[self.timestamp_begin :].set(-float("inf")),
                    scores_k.at[: self.eos_token_id].set(-float("inf")),
                ),
                scores_k,
            )

        scores = jax.vmap(handle_pairs)(input_ids, scores)

        apply_max_initial_timestamp = jnp.where(cur_len == self.begin_index, True, False)
        apply_max_initial_timestamp = jnp.where(
            self.max_initial_timestamp_index is not None,
            True and apply_max_initial_timestamp,
            False,
        )

        last_allowed = self.timestamp_begin + self.max_initial_timestamp_index

        scores = jnp.where(
            apply_max_initial_timestamp,
            scores.at[:, last_allowed + 1 :].set(-float("inf")),
            scores,
        )

        
        logprobs = jax.nn.log_softmax(scores, axis=-1)

        def handle_cumulative_probs(logprobs_k, scores_k):
            timestamp_logprob = jax.nn.logsumexp(logprobs_k[self.timestamp_begin :], axis=-1)
            max_text_token_logprob = jnp.max(logprobs_k[: self.timestamp_begin])
            return jnp.where(
                timestamp_logprob > max_text_token_logprob,
                scores_k.at[: self.timestamp_begin].set(-float("inf")),
                scores_k,
            )

        scores = jax.vmap(handle_cumulative_probs)(logprobs, scores)

        return scores