Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2021 The HuggingFace Inc. team3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16import inspect17 18import jax19import jax.lax as lax20import jax.numpy as jnp21from jax.experimental import sparse22 23from ..utils import add_start_docstrings24from ..utils.logging import get_logger25 26 27logger = get_logger(__name__)28 29 30LOGITS_PROCESSOR_INPUTS_DOCSTRING = r"""31    Args:32        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):33            Indices of input sequence tokens in the vocabulary.34 35            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and36            [`PreTrainedTokenizer.__call__`] for details.37 38            [What are input IDs?](../glossary#input-ids)39        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):40            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam41            search or log softmax for each vocabulary token when using beam search42        kwargs (`dict[str, Any]`, *optional*):43            Additional logits processor specific kwargs.44 45    Return:46        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.47 48"""49 50 51class FlaxLogitsProcessor:52    """Abstract base class for all logit processors that can be applied during generation."""53 54    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)55    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray) -> jnp.ndarray:56        """Flax method for processing logits."""57        raise NotImplementedError(58            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."59        )60 61 62class FlaxLogitsWarper:63    """Abstract base class for all logit warpers that can be applied during generation with multinomial sampling."""64 65    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)66    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray) -> jnp.ndarray:67        """Flax method for warping logits."""68        raise NotImplementedError(69            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."70        )71 72 73class FlaxLogitsProcessorList(list):74    """75    This class can be used to create a list of [`FlaxLogitsProcessor`] or [`FlaxLogitsWarper`] to subsequently process76    a `scores` input tensor. This class inherits from list and adds a specific *__call__* method to apply each77    [`FlaxLogitsProcessor`] or [`FlaxLogitsWarper`] to the inputs.78    """79 80    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)81    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int, **kwargs) -> jnp.ndarray:82        for processor in self:83            function_args = inspect.signature(processor.__call__).parameters84            if len(function_args) > 3:85                if not all(arg in kwargs for arg in list(function_args.keys())[2:]):86                    raise ValueError(87                        f"Make sure that all the required parameters: {list(function_args.keys())} for "88                        f"{processor.__class__} are passed to the logits processor."89                    )90                scores = processor(input_ids, scores, cur_len, **kwargs)91            else:92                scores = processor(input_ids, scores, cur_len)93        return scores94 95 96class FlaxTemperatureLogitsWarper(FlaxLogitsWarper):97    r"""98    [`FlaxLogitsWarper`] for temperature (exponential scaling output probability distribution).99 100    Args:101        temperature (`float`):102            The value used to module the logits distribution.103    """104 105    def __init__(self, temperature: float):106        if not isinstance(temperature, float) or not (temperature > 0):107            raise ValueError(f"`temperature` has to be a strictly positive float, but is {temperature}")108 109        self.temperature = temperature110 111    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:112        scores = scores / self.temperature113        return scores114 115 116class FlaxTopPLogitsWarper(FlaxLogitsWarper):117    """118    [`FlaxLogitsWarper`] that performs top-p, i.e. restricting to top tokens summing to prob_cut_off <= prob_cut_off.119 120    Args:121        top_p (`float`):122            If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or123            higher are kept for generation.124        filter_value (`float`, *optional*, defaults to -inf):125            All filtered values will be set to this float value.126        min_tokens_to_keep (`int`, *optional*, defaults to 1):127            Minimum number of tokens that cannot be filtered.128    """129 130    def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):131        if not isinstance(top_p, float) or (top_p < 0 or top_p > 1.0):132            raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")133        if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):134            raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")135 136        self.top_p = top_p137        self.filter_value = filter_value138        self.min_tokens_to_keep = min_tokens_to_keep139 140    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:141        topk_scores, topk_indices = lax.top_k(scores, scores.shape[-1])142 143        mask_scores = jnp.full_like(scores, self.filter_value)144        cumulative_probs = jax.nn.softmax(topk_scores, axis=-1).cumsum(axis=-1)145        score_mask = cumulative_probs < self.top_p146 147        # include the token that is higher than top_p as well148        score_mask = jnp.roll(score_mask, 1)149        score_mask |= score_mask.at[:, 0].set(True)150 151        # min tokens to keep152        score_mask = score_mask.at[:, : self.min_tokens_to_keep].set(True)153 154        topk_next_scores = jnp.where(score_mask, topk_scores, mask_scores)155        next_scores = jax.lax.sort_key_val(topk_indices, topk_next_scores)[-1]156 157        return next_scores158 159 160class FlaxTopKLogitsWarper(FlaxLogitsWarper):161    r"""162    [`FlaxLogitsWarper`] that performs top-k, i.e. restricting to the k highest probability elements.163 164    Args:165        top_k (`int`):166            The number of highest probability vocabulary tokens to keep for top-k-filtering.167        filter_value (`float`, *optional*, defaults to -inf):168            All filtered values will be set to this float value.169        min_tokens_to_keep (`int`, *optional*, defaults to 1):170            Minimum number of tokens that cannot be filtered.171    """172 173    def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):174        if not isinstance(top_k, int) or top_k <= 0:175            raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}")176 177        self.top_k = max(top_k, min_tokens_to_keep)178        self.filter_value = filter_value179 180    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:181        batch_size, vocab_size = scores.shape182        next_scores_flat = jnp.full(batch_size * vocab_size, self.filter_value)183 184        topk = min(self.top_k, scores.shape[-1])  # Safety check185        topk_scores, topk_indices = lax.top_k(scores, topk)186        shift = jnp.broadcast_to((jnp.arange(batch_size) * vocab_size)[:, None], (batch_size, topk)).flatten()187        topk_scores_flat = topk_scores.flatten()188        topk_indices_flat = topk_indices.flatten() + shift189 190        next_scores_flat = next_scores_flat.at[topk_indices_flat].set(topk_scores_flat)191        next_scores = next_scores_flat.reshape(batch_size, vocab_size)192        return next_scores193 194 195class FlaxForcedBOSTokenLogitsProcessor(FlaxLogitsProcessor):196    r"""197    [`FlaxLogitsProcessor`] that enforces the specified token as the first generated token.198 199    Args:200        bos_token_id (`int`):201            The id of the token to force as the first generated token.202    """203 204    def __init__(self, bos_token_id: int):205        self.bos_token_id = bos_token_id206 207    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:208        new_scores = jnp.full(scores.shape, -float("inf"))209 210        apply_penalty = 1 - jnp.bool_(cur_len - 1)211 212        scores = jnp.where(apply_penalty, new_scores.at[:, self.bos_token_id].set(0), scores)213 214        return scores215 216 217class FlaxForcedEOSTokenLogitsProcessor(FlaxLogitsProcessor):218    r"""219    [`FlaxLogitsProcessor`] that enforces the specified token as the last generated token when `max_length` is reached.220 221    Args:222        max_length (`int`):223            The maximum length of the sequence to be generated.224        eos_token_id (`int`):225            The id of the token to force as the last generated token when `max_length` is reached.226    """227 228    def __init__(self, max_length: int, eos_token_id: int):229        self.max_length = max_length230        self.eos_token_id = eos_token_id231 232    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:233        new_scores = jnp.full(scores.shape, -float("inf"))234 235        apply_penalty = 1 - jnp.bool_(cur_len - self.max_length + 1)236 237        scores = jnp.where(apply_penalty, new_scores.at[:, self.eos_token_id].set(0), scores)238 239        return scores240 241 242class FlaxMinLengthLogitsProcessor(FlaxLogitsProcessor):243    r"""244    [`FlaxLogitsProcessor`] enforcing a min-length by setting EOS probability to 0.245 246    Args:247        min_length (`int`):248            The minimum length below which the score of `eos_token_id` is set to `-float("Inf")`.249        eos_token_id (`int`):250            The id of the *end-of-sequence* token.251    """252 253    def __init__(self, min_length: int, eos_token_id: int):254        if not isinstance(min_length, int) or min_length < 0:255            raise ValueError(f"`min_length` has to be a positive integer, but is {min_length}")256 257        if not isinstance(eos_token_id, int) or eos_token_id < 0:258            raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}")259 260        self.min_length = min_length261        self.eos_token_id = eos_token_id262 263    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:264        # create boolean flag to decide if min length penalty should be applied265        apply_penalty = 1 - jnp.clip(cur_len - self.min_length, 0, 1)266 267        scores = jnp.where(apply_penalty, scores.at[:, self.eos_token_id].set(-float("inf")), scores)268 269        return scores270 271 272class FlaxSuppressTokensAtBeginLogitsProcessor(FlaxLogitsProcessor):273    r"""274    [`FlaxLogitsProcessor`] suppressing a list of tokens as soon as the `generate` function starts generating using275    `begin_index` tokens. This should ensure that the tokens defined by `begin_suppress_tokens` are not sampled at the276    beginning of the generation.277 278    Args:279        begin_suppress_tokens (`list[int]`):280            Tokens to not sample.281        begin_index (`int`):282            Index where the tokens are suppressed.283    """284 285    def __init__(self, begin_suppress_tokens, begin_index):286        self.begin_suppress_tokens = list(begin_suppress_tokens)287        self.begin_index = begin_index288 289    def __call__(self, input_ids, scores, cur_len: int):290        apply_penalty = 1 - jnp.bool_(cur_len - self.begin_index)291 292        scores = jnp.where(apply_penalty, scores.at[:, self.begin_suppress_tokens].set(-float("inf")), scores)293 294        return scores295 296 297class FlaxSuppressTokensLogitsProcessor(FlaxLogitsProcessor):298    r"""299    [`FlaxLogitsProcessor`] suppressing a list of tokens at each decoding step. The processor will set their log probs300    to be `-inf` so they are not sampled.301 302    Args:303        suppress_tokens (`list`):304            Tokens to not sample.305    """306 307    def __init__(self, suppress_tokens: list):308        self.suppress_tokens = list(suppress_tokens)309 310    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:311        scores = scores.at[..., self.suppress_tokens].set(-float("inf"))312 313        return scores314 315 316class FlaxForceTokensLogitsProcessor(FlaxLogitsProcessor):317    r"""318    [`FlaxLogitsProcessor`] that takes a list of pairs of integers which indicates a mapping from generation indices to319    token indices that will be forced before sampling. The processor will set their log probs to 0 and all other tokens320    to `-inf` so that they are sampled at their corresponding index.321 322    Args:323        force_token_map (`list`):324            Map giving token ids and indices where they will be forced to be sampled.325    """326 327    def __init__(self, force_token_map):328        force_token_map = dict(force_token_map)329        # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the330        # index of the array corresponds to the index of the token to be forced, for XLA compatibility.331        # Indexes without forced tokens will have a negative value.332        force_token_array = jnp.ones((max(force_token_map.keys()) + 1), dtype=jnp.int32) * -1333        for index, token in force_token_map.items():334            if token is not None:335                force_token_array = force_token_array.at[index].set(token)336        self.force_token_array = jnp.int32(force_token_array)337 338    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:339        def _force_token(generation_idx):340            batch_size = scores.shape[0]341            current_token = self.force_token_array[generation_idx]342 343            new_scores = jnp.ones_like(scores, dtype=scores.dtype) * -float("inf")344            updates = jnp.zeros((batch_size, 1), dtype=scores.dtype)345            new_scores = lax.dynamic_update_slice(new_scores, updates, (0, current_token))346            return new_scores347 348        scores = lax.cond(349            cur_len >= self.force_token_array.shape[0],350            # If the current length is geq than the length of force_token_array, the processor does nothing.351            lambda: scores,352            # Otherwise, it may force a certain token.353            lambda: lax.cond(354                self.force_token_array[cur_len] >= 0,355                # Only valid (positive) tokens are forced356                lambda: _force_token(cur_len),357                # Otherwise, the processor does nothing.358                lambda: scores,359            ),360        )361        return scores362 363 364class FlaxWhisperTimeStampLogitsProcessor(FlaxLogitsProcessor):365    r"""366    Whisper specific Processor. This processor can be used to force a list of tokens. The processor will set their log367    probs to `inf` so that they are sampled at their corresponding index.368 369    Args:370        generate_config (`GenerateConfig`):371            The generate config used to generate the output. The following parameters are required:372                eos_token_id (`int`, *optional*, defaults to 50257):373                    The id of the *end-of-sequence* token.374                no_timestamps_token_id (`int`, *optional*, defaults to 50363):375                    The id of the `"<|notimestamps|>"` token.376                max_initial_timestamp_index (`int`, *optional*, defaults to 1):377                    Used to set the maximum value of the initial timestamp. This is used to prevent the model from378                    predicting timestamps that are too far in the future.379    """380 381    def __init__(self, generate_config, model_config, decoder_input_length):382        self.eos_token_id = generate_config.eos_token_id383        self.no_timestamps_token_id = generate_config.no_timestamps_token_id384        self.timestamp_begin = generate_config.no_timestamps_token_id + 1385 386        self.begin_index = decoder_input_length + 1387 388        if generate_config.is_multilingual:389            # room for language token and task token390            self.begin_index += 2391        if hasattr(generate_config, "max_initial_timestamp_index"):392            self.max_initial_timestamp_index = generate_config.max_initial_timestamp_index393        else:394            self.max_initial_timestamp_index = model_config.vocab_size395        if self.max_initial_timestamp_index is None:396            self.max_initial_timestamp_index = model_config.vocab_size397 398    def __call__(self, input_ids, scores, cur_len):399        # suppress <|notimestamps|> which is handled by without_timestamps400        scores = scores.at[:, self.no_timestamps_token_id].set(-float("inf"))401 402        def handle_pairs(input_ids_k, scores_k):403            last_was_timestamp = jnp.where((cur_len - self.begin_index) >= 1, True, False)404            last_was_timestamp = jnp.where(405                input_ids_k[cur_len - 1] >= self.timestamp_begin,406                True and last_was_timestamp,407                False,408            )409 410            penultimate_was_timestamp = jnp.where((cur_len - self.begin_index) < 2, True, False)411            penultimate_was_timestamp = jnp.where(412                input_ids_k[cur_len - 2] >= self.timestamp_begin,413                True,414                penultimate_was_timestamp,415            )416 417            return jnp.where(418                last_was_timestamp,419                jnp.where(420                    penultimate_was_timestamp > 0,421                    scores_k.at[self.timestamp_begin :].set(-float("inf")),422                    scores_k.at[: self.eos_token_id].set(-float("inf")),423                ),424                scores_k,425            )426 427        scores = jax.vmap(handle_pairs)(input_ids, scores)428 429        apply_max_initial_timestamp = jnp.where(cur_len == self.begin_index, True, False)430        apply_max_initial_timestamp = jnp.where(431            self.max_initial_timestamp_index is not None,432            True and apply_max_initial_timestamp,433            False,434        )435 436        last_allowed = self.timestamp_begin + self.max_initial_timestamp_index437 438        scores = jnp.where(439            apply_max_initial_timestamp,440            scores.at[:, last_allowed + 1 :].set(-float("inf")),441            scores,442        )443 444        # if sum of probability over timestamps is above any other token, sample timestamp445        logprobs = jax.nn.log_softmax(scores, axis=-1)446 447        def handle_cumulative_probs(logprobs_k, scores_k):448            timestamp_logprob = jax.nn.logsumexp(logprobs_k[self.timestamp_begin :], axis=-1)449            max_text_token_logprob = jnp.max(logprobs_k[: self.timestamp_begin])450            return jnp.where(451                timestamp_logprob > max_text_token_logprob,452                scores_k.at[: self.timestamp_begin].set(-float("inf")),453                scores_k,454            )455 456        scores = jax.vmap(handle_cumulative_probs)(logprobs, scores)457 458        return scores459 460 461class FlaxNoRepeatNGramLogitsProcessor(FlaxLogitsProcessor):462    r"""463    [`FlaxLogitsProcessor`] that enforces no repetition of n-grams. See464    [Fairseq](https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345).465 466    Args:467        ngram_size (`int`):468            All ngrams of size `ngram_size` can only occur once.469    """470 471    def __init__(self, ngram_size: int):472        if not isinstance(ngram_size, int) or ngram_size <= 0:473            raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}")474        self.ngram_size = ngram_size475 476    def get_previous_ngrams(self, input_ids: jnp.ndarray, vocab_size: int, cur_len: int):477        """478        get a matrix of size (batch_size,) + (vocab_size,)*n (for n-grams) that479        represent the n-grams that occurred previously.480        The BCOO representation allow to store only the few non-zero entries, instead of the full (huge) matrix481        """482        batch_size, seq_len = input_ids.shape483        # number of n-grams in the whole sequence484        seq_ngrams = seq_len - (self.ngram_size - 1)485        # number of n-grams in the currently generated sequence486        cur_ngrams = cur_len - (self.ngram_size - 1)487 488        def body_fun(i, val):489            b = i % batch_size490            pos = i // batch_size491            return val.at[i].set(492                jnp.array(493                    [494                        b,495                    ]496                    + [jnp.array(input_ids)[b, pos + j] for j in range(self.ngram_size)]497                )498            )499 500        shape = (batch_size * seq_ngrams, self.ngram_size + 1)501        all_update_indices = jax.lax.fori_loop(502            0, batch_size * cur_ngrams, body_fun, jnp.zeros(shape, dtype=input_ids.dtype)503        )504 505        # ignore the n-grams not yet generated506        data = (jnp.arange(batch_size * seq_ngrams) < batch_size * cur_ngrams).astype("float32")507 508        return sparse.BCOO((data, all_update_indices), shape=(batch_size,) + (vocab_size,) * self.ngram_size)509 510    def get_banned_tokens_mask(self, latest_tokens: jnp.ndarray, previous_ngrams) -> jnp.ndarray:511        """512        Determines which tokens must be banned given latest tokens and the previously seen513        ngrams.514        """515 516        @sparse.sparsify517        @jax.vmap518        def inner_fn(latest_tokens, previous_ngrams):519            return previous_ngrams[tuple(latest_tokens)]520 521        return sparse.bcoo_todense(inner_fn(latest_tokens, previous_ngrams))522 523    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:524        def true_fn():525            _, vocab_size = scores.shape526            # store the previously seen n-grams527            previous_ngrams = self.get_previous_ngrams(input_ids, vocab_size, cur_len)528 529            # get the n-1 last tokens that prefix the n-gram being generated530            latest_tokens = jnp.zeros((input_ids.shape[0], self.ngram_size - 1), dtype=input_ids.dtype)531            latest_tokens = jax.lax.dynamic_update_slice(532                latest_tokens,533                jax.lax.dynamic_slice(534                    input_ids, (0, cur_len - (self.ngram_size - 1)), (input_ids.shape[0], (self.ngram_size - 1))535                ),536                (0, 0),537            )538 539            # compute the banned tokens, ie all the tokens that when added to the latest tokens lead to a n-gram that was previously generated540            banned_tokens_indices_mask = self.get_banned_tokens_mask(latest_tokens, previous_ngrams).astype("bool")541            return jnp.where(banned_tokens_indices_mask, -float("inf"), scores)542 543        output = jax.lax.cond((cur_len >= self.ngram_size - 1), true_fn, lambda: scores)544        return output545