
"""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


lowercase_         =       get_logger(__name__)


lowercase_         =       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_						:
 """simple docstring"""




 @add_start_docstrings(a_			)
 def __call__(			self							:			Optional[Any]  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray			)->       jnp.ndarray:



      """simple docstring"""

      raise NotImplementedError(
          f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''			)



class 		UpperCAmelCase_						:
 """simple docstring"""




 @add_start_docstrings(a_			)
 def __call__(			self							:			Tuple  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray			)->       jnp.ndarray:



      """simple docstring"""

      raise NotImplementedError(
          f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''			)



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 @add_start_docstrings(a_			)
 def __call__(			self							:			List[str]  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int  , **a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      for processor in self:
           UpperCAmelCase_    :     List[str]            =	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_    :     Dict            =	processor(a_  , a_  , a_  , **a_			)
           else:
                UpperCAmelCase_    :     int            =	processor(a_  , a_  , a_			)
      return scores



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Dict  , a_							:			float			)->       int:



      """simple docstring"""

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

      UpperCAmelCase_    :     int            =	temperature






 def __call__(			self							:			Union[str, Any]  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[Any]            =	scores / self.temperature
      return scores



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Optional[Any]  , a_							:			float  , a_							:			float = -float("""Inf"""			)  , a_							:			int = 1			)->       Dict:



      """simple docstring"""

      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[Any]            =	top_p
      UpperCAmelCase_    :     Union[str, Any]            =	filter_value
      UpperCAmelCase_    :     List[str]            =	min_tokens_to_keep






 def __call__(			self							:			Union[str, Any]  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      UpperCAmelCase_       ,UpperCAmelCase_    :     Dict            =	lax.top_k(a_  , scores.shape[-1]			)

      UpperCAmelCase_    :     str            =	jnp.full_like(a_  , self.filter_value			)
      UpperCAmelCase_    :     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_    :     Optional[int]            =	jnp.roll(a_  , 1			)
      score_mask |= score_mask.at[:, 0].set(a_			)

      # min tokens to keep
      UpperCAmelCase_    :     Tuple            =	score_mask.at[:, : self.min_tokens_to_keep].set(a_			)

      UpperCAmelCase_    :     Dict            =	jnp.where(a_  , a_  , a_			)
      UpperCAmelCase_    :     List[str]            =	jax.lax.sort_key_val(a_  , a_			)[-1]

      return next_scores



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Any  , a_							:			int  , a_							:			float = -float("""Inf"""			)  , a_							:			int = 1			)->       Optional[int]:



      """simple docstring"""

      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_    :     Any            =	max(a_  , a_			)
      UpperCAmelCase_    :     str            =	filter_value






 def __call__(			self							:			Optional[int]  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      UpperCAmelCase_       ,UpperCAmelCase_    :     Optional[Any]            =	scores.shape
      UpperCAmelCase_    :     Optional[Any]            =	jnp.full(batch_size * vocab_size  , self.filter_value			)

      UpperCAmelCase_    :     Tuple            =	min(self.top_k  , scores.shape[-1]			)  # Safety check
      UpperCAmelCase_       ,UpperCAmelCase_    :     Any            =	lax.top_k(a_  , a_			)
      UpperCAmelCase_    :     List[str]            =	jnp.broadcast_to((jnp.arange(a_			) * vocab_size)[:, None]  , (batch_size, topk)			).flatten()
      UpperCAmelCase_    :     Optional[Any]            =	topk_scores.flatten()
      UpperCAmelCase_    :     Tuple            =	topk_indices.flatten() + shift

      UpperCAmelCase_    :     Union[str, Any]            =	next_scores_flat.at[topk_indices_flat].set(a_			)
      UpperCAmelCase_    :     str            =	next_scores_flat.reshape(a_  , a_			)
      return next_scores



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Tuple  , a_							:			int			)->       Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Tuple            =	bos_token_id






 def __call__(			self							:			str  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      UpperCAmelCase_    :     Tuple            =	jnp.full(scores.shape  , -float("""inf"""			)			)

      UpperCAmelCase_    :     List[Any]            =	1 - jnp.bool_(cur_len - 1			)

      UpperCAmelCase_    :     List[str]            =	jnp.where(a_  , new_scores.at[:, self.bos_token_id].set(0			)  , a_			)

      return scores



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Tuple  , a_							:			int  , a_							:			int			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     List[str]            =	max_length
      UpperCAmelCase_    :     Dict            =	eos_token_id






 def __call__(			self							:			List[Any]  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      UpperCAmelCase_    :     int            =	jnp.full(scores.shape  , -float("""inf"""			)			)

      UpperCAmelCase_    :     Optional[Any]            =	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_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			List[Any]  , a_							:			int  , a_							:			int			)->       Optional[int]:



      """simple docstring"""

      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_    :     str            =	min_length
      UpperCAmelCase_    :     int            =	eos_token_id






 def __call__(			self							:			Union[str, Any]  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      UpperCAmelCase_    :     int            =	1 - jnp.clip(cur_len - self.min_length  , 0  , 1			)

      UpperCAmelCase_    :     List[str]            =	jnp.where(a_  , scores.at[:, self.eos_token_id].set(-float("""inf"""			)			)  , a_			)

      return scores



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Union[str, Any]  , a_							:			List[str]  , a_							:			Optional[int]			)->       str:



      """simple docstring"""

      UpperCAmelCase_    :     List[str]            =	list(a_			)
      UpperCAmelCase_    :     List[Any]            =	begin_index






 def __call__(			self							:			str  , a_							:			Union[str, Any]  , a_							:			Any  , a_							:			int			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     List[str]            =	1 - jnp.bool_(cur_len - self.begin_index			)

      UpperCAmelCase_    :     Optional[int]            =	jnp.where(a_  , scores.at[:, self.begin_suppress_tokens].set(-float("""inf"""			)			)  , a_			)

      return scores



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Dict  , a_							:			list			)->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     str            =	list(a_			)






 def __call__(			self							:			Any  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      UpperCAmelCase_    :     Tuple            =	scores.at[..., self.suppress_tokens].set(-float("""inf"""			)			)

      return scores



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			List[str]  , a_							:			Dict			)->       Optional[int]:



      """simple docstring"""

      UpperCAmelCase_    :     str            =	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_    :     Union[str, 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_    :     Dict            =	force_token_array.at[index].set(a_			)
      UpperCAmelCase_    :     List[Any]            =	jnp.intaa(a_			)






 def __call__(			self							:			Optional[Any]  , a_							:			jnp.ndarray  , a_							:			jnp.ndarray  , a_							:			int			)->       jnp.ndarray:



      """simple docstring"""

      def _force_token(a_							:			List[str]			):
           UpperCAmelCase_    :     Dict            =	scores.shape[0]
           UpperCAmelCase_    :     Dict            =	self.force_token_array[generation_idx]

           UpperCAmelCase_    :     str            =	jnp.ones_like(a_  , dtype=scores.dtype			) * -float("""inf"""			)
           UpperCAmelCase_    :     Tuple            =	jnp.zeros((batch_size, 1)  , dtype=scores.dtype			)
           UpperCAmelCase_    :     Tuple            =	lax.dynamic_update_slice(a_  , a_  , (0, current_token)			)
           return new_scores

      UpperCAmelCase_    :     List[str]            =	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_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			Dict  , a_							:			Tuple  , a_							:			Optional[Any]  , a_							:			Optional[int]			)->       Dict:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	generate_config.eos_token_id
      UpperCAmelCase_    :     str            =	generate_config.no_timestamps_token_id
      UpperCAmelCase_    :     int            =	generate_config.no_timestamps_token_id + 1

      UpperCAmelCase_    :     Dict            =	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_    :     str            =	generate_config.max_initial_timestamp_index
      else:
           UpperCAmelCase_    :     List[str]            =	model_config.vocab_size
      if self.max_initial_timestamp_index is None:
           UpperCAmelCase_    :     int            =	model_config.vocab_size






 def __call__(			self							:			str  , a_							:			Dict  , a_							:			Tuple  , a_							:			int			)->       Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[Any]            =	scores.at[:, self.no_timestamps_token_id].set(-float("""inf"""			)			)

      def handle_pairs(a_							:			Optional[Any]  , a_							:			List[Any]			):
           UpperCAmelCase_    :     Any            =	jnp.where((cur_len - self.begin_index) >= 1  , a_  , a_			)
           UpperCAmelCase_    :     Optional[int]            =	jnp.where(
               input_ids_k[cur_len - 1] >= self.timestamp_begin  , True and last_was_timestamp  , a_  , )

           UpperCAmelCase_    :     Optional[int]            =	jnp.where((cur_len - self.begin_index) < 2  , a_  , a_			)
           UpperCAmelCase_    :     Tuple            =	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_    :     str            =	jax.vmap(a_			)(a_  , a_			)

      UpperCAmelCase_    :     Tuple            =	jnp.where(cur_len == self.begin_index  , a_  , a_			)
      UpperCAmelCase_    :     Any            =	jnp.where(
          self.max_initial_timestamp_index is not None  , True and apply_max_initial_timestamp  , a_  , )

      UpperCAmelCase_    :     List[str]            =	self.timestamp_begin + self.max_initial_timestamp_index

      UpperCAmelCase_    :     Dict            =	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_    :     Tuple            =	jax.nn.log_softmax(a_  , axis=-1			)

      def handle_cumulative_probs(a_							:			str  , a_							:			Any			):
           UpperCAmelCase_    :     str            =	jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]  , axis=-1			)
           UpperCAmelCase_    :     List[str]            =	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_    :     List[str]            =	jax.vmap(a_			)(a_  , a_			)

      return scores



"""simple docstring"""







import json
from typing import Iterator, List, Union

from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, trainers
from tokenizers.implementations.base_tokenizer import BaseTokenizer
from tokenizers.models import Unigram
from tokenizers.processors import TemplateProcessing



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			List[Any]  , a_							:			str = "▁"  , a_							:			bool = True  , a_							:			Union[str, AddedToken] = "<unk>"  , a_							:			Union[str, AddedToken] = "</s>"  , a_							:			Union[str, AddedToken] = "<pad>"  , )->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	{
          """pad""": {"""id""": 0, """token""": pad_token},
          """eos""": {"""id""": 1, """token""": eos_token},
          """unk""": {"""id""": 2, """token""": unk_token},
      }

      UpperCAmelCase_    :     Dict            =	[None] * len(self.special_tokens			)
      for token_dict in self.special_tokens.values():
           UpperCAmelCase_    :     Optional[int]            =	token_dict["""token"""]

      UpperCAmelCase_    :     Tuple            =	Tokenizer(Unigram()			)

      UpperCAmelCase_    :     Optional[Any]            =	normalizers.Sequence(
          [
              normalizers.Nmt(),
              normalizers.NFKC(),
              normalizers.Replace(Regex(""" {2,}"""			)  , """ """			),
              normalizers.Lowercase(),
          ]			)
      UpperCAmelCase_    :     Any            =	pre_tokenizers.Sequence(
          [
              pre_tokenizers.Metaspace(replacement=a_  , add_prefix_space=a_			),
              pre_tokenizers.Digits(individual_digits=a_			),
              pre_tokenizers.Punctuation(),
          ]			)
      UpperCAmelCase_    :     str            =	decoders.Metaspace(replacement=a_  , add_prefix_space=a_			)

      UpperCAmelCase_    :     List[Any]            =	TemplateProcessing(
          single=f'''$A {self.special_tokens['eos']['token']}'''  , special_tokens=[(self.special_tokens["""eos"""]["""token"""], self.special_tokens["""eos"""]["""id"""])]  , )

      UpperCAmelCase_    :     Dict            =	{
          """model""": """SentencePieceUnigram""",
          """replacement""": replacement,
          """add_prefix_space""": add_prefix_space,
      }

      super().__init__(a_  , a_			)




 def 	a   (			self							:			int  , a_							:			Union[str, List[str]]  , a_							:			int = 80_00  , a_							:			bool = True  , )->       int:



      """simple docstring"""

      UpperCAmelCase_    :     int            =	trainers.UnigramTrainer(
          vocab_size=a_  , special_tokens=self.special_tokens_list  , show_progress=a_  , )

      if isinstance(a_  , a_			):
           UpperCAmelCase_    :     str            =	[files]
      self._tokenizer.train(a_  , trainer=a_			)

      self.add_unk_id()




 def 	a   (			self							:			List[Any]  , a_							:			Union[Iterator[str], Iterator[Iterator[str]]]  , a_							:			int = 80_00  , a_							:			bool = True  , )->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     Union[str, Any]            =	trainers.UnigramTrainer(
          vocab_size=a_  , special_tokens=self.special_tokens_list  , show_progress=a_  , )

      self._tokenizer.train_from_iterator(a_  , trainer=a_			)

      self.add_unk_id()






 def 	a   (			self							:			Dict			)->       Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_    :     int            =	json.loads(self._tokenizer.to_str()			)

      UpperCAmelCase_    :     Any            =	self.special_tokens["""unk"""]["""id"""]

      UpperCAmelCase_    :     Tuple            =	Tokenizer.from_str(json.dumps(a_			)			)

