



"""simple docstring"""




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


_lowerCAmelCase							:   int				=				get_logger(__name__)


_lowerCAmelCase							:   str				=				r"\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n"






class       UpperCAmelCase_    :



      @add_start_docstrings(A		)
      def __call__(	self	:      int	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray		):
           raise NotImplementedError(
               f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'		)







class       UpperCAmelCase_    :



      @add_start_docstrings(A		)
      def __call__(	self	:      List[Any]	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray		):
           raise NotImplementedError(
               f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'		)







class       UpperCAmelCase_    (							_UpperCamelCase		):



      @add_start_docstrings(A		)
      def __call__(	self	:      Tuple	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int	,      **A	:      Optional[Any]		):
           for processor in self:
                _UpperCAmelCase     :      Any														=       inspect.signature(processor.__call__		).parameters
                if len(A		) > 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.'		)
                     _UpperCAmelCase     :      Optional[int]														=       processor(A	,      A	,      A	,      **A		)
                else:
                     _UpperCAmelCase     :      Dict														=       processor(A	,      A	,      A		)
           return scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



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

           _UpperCAmelCase     :      List[str]														=       temperature



      def __call__(	self	:      int	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int		):
           _UpperCAmelCase     :      Optional[Any]														=       scores / self.temperature
           return scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



      def __init__(	self	:      Tuple	,      A	:      float	,      A	:      float = -float("Inf"		)	,      A	:      int = 1		):
           if not isinstance(A	,      A		) 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(A	,      A		) 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}'		)

           _UpperCAmelCase     :      List[str]														=       top_p
           _UpperCAmelCase     :      List[Any]														=       filter_value
           _UpperCAmelCase     :      Optional[Any]														=       min_tokens_to_keep



      def __call__(	self	:      int	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      Optional[int]														=       lax.top_k(A	,      scores.shape[-1]		)

           _UpperCAmelCase     :      str														=       jnp.full_like(A	,      self.filter_value		)
           _UpperCAmelCase     :      Optional[int]														=       jax.nn.softmax(A	,      axis=-1		).cumsum(axis=-1		)
           _UpperCAmelCase     :      List[Any]														=       cumulative_probs < self.top_p

           # include the token that is higher than top_p as well
           _UpperCAmelCase     :      List[Any]														=       jnp.roll(A	,      1		)
           score_mask |= score_mask.at[:, 0].set(A		)

           # min tokens to keep
           _UpperCAmelCase     :      Optional[int]														=       score_mask.at[:, : self.min_tokens_to_keep].set(A		)

           _UpperCAmelCase     :      int														=       jnp.where(A	,      A	,      A		)
           _UpperCAmelCase     :      List[Any]														=       jax.lax.sort_key_val(A	,      A		)[-1]

           return next_scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



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

           _UpperCAmelCase     :      str														=       max(A	,      A		)
           _UpperCAmelCase     :      List[Any]														=       filter_value



      def __call__(	self	:      str	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      List[str]														=       scores.shape
           _UpperCAmelCase     :      List[str]														=       jnp.full(batch_size * vocab_size	,      self.filter_value		)

           _UpperCAmelCase     :      Dict														=       min(self.top_k	,      scores.shape[-1]		)  # Safety check
           _UpperCAmelCase							,			_UpperCAmelCase     :      int														=       lax.top_k(A	,      A		)
           _UpperCAmelCase     :      Tuple														=       jnp.broadcast_to((jnp.arange(A		) * vocab_size)[:, None]	,      (batch_size, topk)		).flatten()
           _UpperCAmelCase     :      str														=       topk_scores.flatten()
           _UpperCAmelCase     :      int														=       topk_indices.flatten() + shift

           _UpperCAmelCase     :      Union[str, Any]														=       next_scores_flat.at[topk_indices_flat].set(A		)
           _UpperCAmelCase     :      Dict														=       next_scores_flat.reshape(A	,      A		)
           return next_scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



      def __init__(	self	:      List[Any]	,      A	:      int		):
           _UpperCAmelCase     :      Optional[int]														=       bos_token_id



      def __call__(	self	:      List[str]	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int		):
           _UpperCAmelCase     :      List[str]														=       jnp.full(scores.shape	,      -float("inf"		)		)

           _UpperCAmelCase     :      List[str]														=       1 - jnp.bool_(cur_len - 1		)

           _UpperCAmelCase     :      Dict														=       jnp.where(A	,      new_scores.at[:, self.bos_token_id].set(0		)	,      A		)

           return scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



      def __init__(	self	:      Optional[Any]	,      A	:      int	,      A	:      int		):
           _UpperCAmelCase     :      List[Any]														=       max_length
           _UpperCAmelCase     :      Optional[Any]														=       eos_token_id



      def __call__(	self	:      Union[str, Any]	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int		):
           _UpperCAmelCase     :      Optional[int]														=       jnp.full(scores.shape	,      -float("inf"		)		)

           _UpperCAmelCase     :      List[str]														=       1 - jnp.bool_(cur_len - self.max_length + 1		)

           _UpperCAmelCase     :      List[str]														=       jnp.where(A	,      new_scores.at[:, self.eos_token_id].set(0		)	,      A		)

           return scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



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

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

           _UpperCAmelCase     :      Dict														=       min_length
           _UpperCAmelCase     :      Union[str, Any]														=       eos_token_id



      def __call__(	self	:      Any	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int		):
           # create boolean flag to decide if min length penalty should be applied
           _UpperCAmelCase     :      List[str]														=       1 - jnp.clip(cur_len - self.min_length	,      0	,      1		)

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

           return scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



      def __init__(	self	:      str	,      A	:      Dict	,      A	:      List[Any]		):
           _UpperCAmelCase     :      Any														=       list(A		)
           _UpperCAmelCase     :      str														=       begin_index



      def __call__(	self	:      List[str]	,      A	:      Optional[int]	,      A	:      Optional[Any]	,      A	:      int		):
           _UpperCAmelCase     :      Optional[int]														=       1 - jnp.bool_(cur_len - self.begin_index		)

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

           return scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



      def __init__(	self	:      List[Any]	,      A	:      list		):
           _UpperCAmelCase     :      Union[str, Any]														=       list(A		)



      def __call__(	self	:      Union[str, Any]	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int		):
           _UpperCAmelCase     :      int														=       scores.at[..., self.suppress_tokens].set(-float("inf"		)		)

           return scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



      def __init__(	self	:      Any	,      A	:      Union[str, Any]		):
           _UpperCAmelCase     :      int														=       dict(A		)
           # 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.
           _UpperCAmelCase     :      Any														=       jnp.ones((max(force_token_map.keys()		) + 1)	,      dtype=jnp.intaa		) * -1
           for index, token in force_token_map.items():
                if token is not None:
                     _UpperCAmelCase     :      Optional[Any]														=       force_token_array.at[index].set(A		)
           _UpperCAmelCase     :      Optional[Any]														=       jnp.intaa(A		)



      def __call__(	self	:      List[Any]	,      A	:      jnp.ndarray	,      A	:      jnp.ndarray	,      A	:      int		):
           def _force_token(A	:      List[str]		):
                _UpperCAmelCase     :      int														=       scores.shape[0]
                _UpperCAmelCase     :      List[Any]														=       self.force_token_array[generation_idx]

                _UpperCAmelCase     :      str														=       jnp.ones_like(A	,      dtype=scores.dtype		) * -float("inf"		)
                _UpperCAmelCase     :      Dict														=       jnp.zeros((batch_size, 1)	,      dtype=scores.dtype		)
                _UpperCAmelCase     :      Optional[int]														=       lax.dynamic_update_slice(A	,      A	,      (0, current_token)		)
                return new_scores

           _UpperCAmelCase     :      Optional[int]														=       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(A		)	,      lambda: scores	,      )	,      )
           return scores







class       UpperCAmelCase_    (							_UpperCamelCase		):



      def __init__(	self	:      Any	,      A	:      int	,      A	:      Tuple	,      A	:      Optional[int]		):
           _UpperCAmelCase     :      List[str]														=       generate_config.eos_token_id
           _UpperCAmelCase     :      List[Any]														=       generate_config.no_timestamps_token_id
           _UpperCAmelCase     :      Any														=       generate_config.no_timestamps_token_id + 1

           _UpperCAmelCase     :      Optional[int]														=       decoder_input_length + 1

           if generate_config.is_multilingual:
                # room for language token and task token
                self.begin_index += 2
           if hasattr(A	,      "max_initial_timestamp_index"		):
                _UpperCAmelCase     :      Any														=       generate_config.max_initial_timestamp_index
           else:
                _UpperCAmelCase     :      str														=       model_config.vocab_size
           if self.max_initial_timestamp_index is None:
                _UpperCAmelCase     :      Optional[int]														=       model_config.vocab_size



      def __call__(	self	:      int	,      A	:      Union[str, Any]	,      A	:      int	,      A	:      Optional[int]		):
           # suppress <|notimestamps|> which is handled by without_timestamps
           _UpperCAmelCase     :      Tuple														=       scores.at[:, self.no_timestamps_token_id].set(-float("inf"		)		)

           def handle_pairs(A	:      List[str]	,      A	:      Tuple		):
                _UpperCAmelCase     :      str														=       jnp.where((cur_len - self.begin_index) >= 1	,      A	,      A		)
                _UpperCAmelCase     :      Any														=       jnp.where(
                    input_ids_k[cur_len - 1] >= self.timestamp_begin	,      True and last_was_timestamp	,      A	,      )

                _UpperCAmelCase     :      int														=       jnp.where((cur_len - self.begin_index) < 2	,      A	,      A		)
                _UpperCAmelCase     :      List[Any]														=       jnp.where(
                    input_ids_k[cur_len - 2] >= self.timestamp_begin	,      A	,      A	,      )

                return jnp.where(
                    A	,      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"		)		)	,      )	,      A	,      )

           _UpperCAmelCase     :      Tuple														=       jax.vmap(A		)(A	,      A		)

           _UpperCAmelCase     :      Tuple														=       jnp.where(cur_len == self.begin_index	,      A	,      A		)
           _UpperCAmelCase     :      str														=       jnp.where(
               self.max_initial_timestamp_index is not None	,      True and apply_max_initial_timestamp	,      A	,      )

           _UpperCAmelCase     :      Union[str, Any]														=       self.timestamp_begin + self.max_initial_timestamp_index

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

           # if sum of probability over timestamps is above any other token, sample timestamp
           _UpperCAmelCase     :      Any														=       jax.nn.log_softmax(A	,      axis=-1		)

           def handle_cumulative_probs(A	:      Optional[Any]	,      A	:      int		):
                _UpperCAmelCase     :      Union[str, Any]														=       jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]	,      axis=-1		)
                _UpperCAmelCase     :      Union[str, Any]														=       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"		)		)	,      A	,      )

           _UpperCAmelCase     :      int														=       jax.vmap(A		)(A	,      A		)

           return scores





"""simple docstring"""




import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_gpta import GPTaTokenizer


if TYPE_CHECKING:
 from transformers.pipelines.conversational import Conversation


_lowerCAmelCase							:   Any				=				logging.get_logger(__name__)

_lowerCAmelCase							:   Any				=				{"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

_lowerCAmelCase							:   List[str]				=				{
    "vocab_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/vocab.json",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/vocab.json",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/vocab.json",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/vocab.json",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/vocab.json",
    },
    "merges_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/merges.txt",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/merges.txt",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/merges.txt",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/merges.txt",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/merges.txt",
    },
    "tokenizer_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/tokenizer.json",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/tokenizer.json",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/tokenizer.json",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/tokenizer.json",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/tokenizer.json",
    },
}

_lowerCAmelCase							:   List[str]				=				{
    "gpt2": 10_24,
    "gpt2-medium": 10_24,
    "gpt2-large": 10_24,
    "gpt2-xl": 10_24,
    "distilgpt2": 10_24,
}






class       UpperCAmelCase_    (							_UpperCamelCase		):
      __SCREAMING_SNAKE_CASE		:  Optional[Any]						=					VOCAB_FILES_NAMES
      __SCREAMING_SNAKE_CASE		:  int						=					PRETRAINED_VOCAB_FILES_MAP
      __SCREAMING_SNAKE_CASE		:  Optional[int]						=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __SCREAMING_SNAKE_CASE		:  Optional[int]						=					['input_ids', 'attention_mask']
      __SCREAMING_SNAKE_CASE		:  Dict						=					GPTaTokenizer



      def __init__(	self	:      Optional[Any]	,      A	:      Optional[int]=None	,      A	:      List[Any]=None	,      A	:      Optional[int]=None	,      A	:      Optional[Any]="<|endoftext|>"	,      A	:      Union[str, Any]="<|endoftext|>"	,      A	:      Optional[Any]="<|endoftext|>"	,      A	:      Optional[Any]=False	,      **A	:      str	,      ):
           super().__init__(
               A	,      A	,      tokenizer_file=A	,      unk_token=A	,      bos_token=A	,      eos_token=A	,      add_prefix_space=A	,      **A	,      )

           _UpperCAmelCase     :      List[Any]														=       kwargs.pop("add_bos_token"	,      A		)

           _UpperCAmelCase     :      List[str]														=       json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()		)
           if pre_tok_state.get("add_prefix_space"	,      A		) != add_prefix_space:
                _UpperCAmelCase     :      Dict														=       getattr(A	,      pre_tok_state.pop("type"		)		)
                _UpperCAmelCase     :      Optional[Any]														=       add_prefix_space
                _UpperCAmelCase     :      Optional[Any]														=       pre_tok_class(**A		)

           _UpperCAmelCase     :      List[Any]														=       add_prefix_space



      def 		snake_case_   (	self	:      List[Any]	,      *A	:      List[Any]	,      **A	:      List[str]		):
           _UpperCAmelCase     :      Union[str, Any]														=       kwargs.get("is_split_into_words"	,      A		)
           assert self.add_prefix_space or not is_split_into_words, (
               f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
               "to use it with pretokenized inputs."
           )

           return super()._batch_encode_plus(*A	,      **A		)



      def 		snake_case_   (	self	:      Optional[Any]	,      *A	:      List[Any]	,      **A	:      int		):
           _UpperCAmelCase     :      Union[str, Any]														=       kwargs.get("is_split_into_words"	,      A		)

           assert self.add_prefix_space or not is_split_into_words, (
               f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
               "to use it with pretokenized inputs."
           )

           return super()._encode_plus(*A	,      **A		)



      def 		snake_case_   (	self	:      Optional[int]	,      A	:      str	,      A	:      Optional[str] = None		):
           _UpperCAmelCase     :      List[Any]														=       self._tokenizer.model.save(A	,      name=A		)
           return tuple(A		)



      def 		snake_case_   (	self	:      int	,      A	:      "Conversation"		):
           _UpperCAmelCase     :      List[str]														=       []
           for is_user, text in conversation.iter_texts():
                input_ids.extend(self.encode(A	,      add_special_tokens=A		) + [self.eos_token_id]		)

           if len(A		) > self.model_max_length:
                _UpperCAmelCase     :      Optional[int]														=       input_ids[-self.model_max_length :]
           return input_ids
