

import datasets

from .nmt_bleu import compute_bleu  # From: https://github.com/tensorflow/nmt/blob/master/nmt/scripts/bleu.py


lowerCAmelCase   :      Union[str, Any]     =     '\\n@INPROCEEDINGS{Papineni02bleu:a,\n    author = {Kishore Papineni and Salim Roukos and Todd Ward and Wei-jing Zhu},\n    title = {BLEU: a Method for Automatic Evaluation of Machine Translation},\n    booktitle = {},\n    year = {2002},\n    pages = {311--318}\n}\n@inproceedings{lin-och-2004-orange,\n    title = "{ORANGE}: a Method for Evaluating Automatic Evaluation Metrics for Machine Translation",\n    author = "Lin, Chin-Yew  and\n      Och, Franz Josef",\n    booktitle = "{COLING} 2004: Proceedings of the 20th International Conference on Computational Linguistics",\n    month = "aug 23{--}aug 27",\n    year = "2004",\n    address = "Geneva, Switzerland",\n    publisher = "COLING",\n    url = "https://www.aclweb.org/anthology/C04-1072",\n    pages = "501--507",\n}\n'

lowerCAmelCase   :      Optional[Any]     =     '\\nBLEU (bilingual evaluation understudy) is an algorithm for evaluating the quality of text which has been machine-translated from one natural language to another.\nQuality is considered to be the correspondence between a machine\'s output and that of a human: "the closer a machine translation is to a professional human translation,\nthe better it is" – this is the central idea behind BLEU. BLEU was one of the first metrics to claim a high correlation with human judgements of quality, and\nremains one of the most popular automated and inexpensive metrics.\n\nScores are calculated for individual translated segments—generally sentences—by comparing them with a set of good quality reference translations.\nThose scores are then averaged over the whole corpus to reach an estimate of the translation\'s overall quality. Intelligibility or grammatical correctness\nare not taken into account[citation needed].\n\nBLEU\'s output is always a number between 0 and 1. This value indicates how similar the candidate text is to the reference texts, with values closer to 1\nrepresenting more similar texts. Few human translations will attain a score of 1, since this would indicate that the candidate is identical to one of the\nreference translations. For this reason, it is not necessary to attain a score of 1. Because there are more opportunities to match, adding additional\nreference translations will increase the BLEU score.\n'

lowerCAmelCase   :      int     =     '\nComputes BLEU score of translated segments against one or more references.\nArgs:\n    predictions: list of translations to score.\n        Each translation should be tokenized into a list of tokens.\n    references: list of lists of references for each translation.\n        Each reference should be tokenized into a list of tokens.\n    max_order: Maximum n-gram order to use when computing BLEU score.\n    smooth: Whether or not to apply Lin et al. 2004 smoothing.\nReturns:\n    \'bleu\': bleu score,\n    \'precisions\': geometric mean of n-gram precisions,\n    \'brevity_penalty\': brevity penalty,\n    \'length_ratio\': ratio of lengths,\n    \'translation_length\': translation_length,\n    \'reference_length\': reference_length\nExamples:\n\n    >>> predictions = [\n    ...     ["hello", "there", "general", "kenobi"],                             # tokenized prediction of the first sample\n    ...     ["foo", "bar", "foobar"]                                             # tokenized prediction of the second sample\n    ... ]\n    >>> references = [\n    ...     [["hello", "there", "general", "kenobi"], ["hello", "there", "!"]],  # tokenized references for the first sample (2 references)\n    ...     [["foo", "bar", "foobar"]]                                           # tokenized references for the second sample (1 reference)\n    ... ]\n    >>> bleu = datasets.load_metric("bleu")\n    >>> results = bleu.compute(predictions=predictions, references=references)\n    >>> print(results["bleu"])\n    1.0\n'
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION      ,   _KWARGS_DESCRIPTION)
class        _A						(     datasets.Metric):

  def        UpperCAmelCase			( self    ):





        """simple docstring"""







        return datasets.MetricInfo(
            description=_DESCRIPTION       ,    citation=_CITATION       ,    inputs_description=_KWARGS_DESCRIPTION       ,    features=datasets.Features(
                {
                    'predictions': datasets.Sequence(datasets.Value('string'       ,    id='token'    )       ,    id='sequence'    ),
                    'references': datasets.Sequence(
                        datasets.Sequence(datasets.Value('string'       ,    id='token'    )       ,    id='sequence'    )       ,    id='references'    ),
                }    )       ,    codebase_urls=['https://github.com/tensorflow/nmt/blob/master/nmt/scripts/bleu.py']       ,    reference_urls=[
                'https://en.wikipedia.org/wiki/BLEU',
                'https://towardsdatascience.com/evaluating-text-output-in-nlp-bleu-at-your-own-risk-e8609665a213',
            ]       ,    )



  def        UpperCAmelCase			( self       ,    _SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE=4       ,    _SCREAMING_SNAKE_CASE=False    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     List[str]							=   compute_bleu(
            reference_corpus=snake_case_       ,    translation_corpus=snake_case_       ,    max_order=snake_case_       ,    smooth=snake_case_    )
        (SCREAMING_SNAKE_CASE_) :     int							=   score
        return {
            "bleu": bleu,
            "precisions": precisions,
            "brevity_penalty": bp,
            "length_ratio": ratio,
            "translation_length": translation_length,
            "reference_length": reference_length,
        }













"""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_     :						Dict						=					get_logger(__name__)


lowerCamelCase_     :						List[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  :



      '''simple docstring'''

      @add_start_docstrings(snake_case_  )
      def __call__(     self					,      snake_case_					,      snake_case_  ):

        """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(snake_case_  )
      def __call__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      @add_start_docstrings(snake_case_  )
      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      **snake_case_  ):

        """simple docstring"""

        for processor in self:
          A_  :       Tuple										=  inspect.signature(processor.__call__  ).parameters
          if len(snake_case_  ) > 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."""  )
            A_  :       Tuple										=  processor(snake_case_					,      snake_case_					,      snake_case_					,      **snake_case_  )
          else:
            A_  :       Optional[Any]										=  processor(snake_case_					,      snake_case_					,      snake_case_  )
        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

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

        A_  :       Optional[int]										=  temperature

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       int										=  scores / self.temperature
        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_ = -float('Inf'  )					,      snake_case_ = 1  ):

        """simple docstring"""

        if not isinstance(snake_case_					,      snake_case_  ) 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(snake_case_					,      snake_case_  ) 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}"""  )

        A_  :       str										=  top_p
        A_  :       Union[str, Any]										=  filter_value
        A_  :       int										=  min_tokens_to_keep

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_						,       A_  :       Tuple										=  lax.top_k(snake_case_					,      scores.shape[-1]  )

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

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

        # min tokens to keep
        A_  :       int										=  score_mask.at[:, : self.min_tokens_to_keep].set(snake_case_  )

        A_  :       Optional[Any]										=  jnp.where(snake_case_					,      snake_case_					,      snake_case_  )
        A_  :       List[Any]										=  jax.lax.sort_key_val(snake_case_					,      snake_case_  )[-1]

        return next_scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_ = -float('Inf'  )					,      snake_case_ = 1  ):

        """simple docstring"""

        if not isinstance(snake_case_					,      snake_case_  ) or top_k <= 0:
          raise ValueError(F"""`top_k` has to be a strictly positive integer, but is {top_k}"""  )

        A_  :       str										=  max(snake_case_					,      snake_case_  )
        A_  :       Union[str, Any]										=  filter_value

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_						,       A_  :       int										=  scores.shape
        A_  :       Tuple										=  jnp.full(batch_size * vocab_size					,      self.filter_value  )

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

        A_  :       List[str]										=  next_scores_flat.at[topk_indices_flat].set(snake_case_  )
        A_  :       Union[str, Any]										=  next_scores_flat.reshape(snake_case_					,      snake_case_  )
        return next_scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  bos_token_id

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  jnp.full(scores.shape					,      -float('inf'  )  )

        A_  :       Union[str, Any]										=  1 - jnp.bool_(cur_len - 1  )

        A_  :       str										=  jnp.where(snake_case_					,      new_scores.at[:, self.bos_token_id].set(0  )					,      snake_case_  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Dict										=  max_length
        A_  :       Optional[int]										=  eos_token_id

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  jnp.full(scores.shape					,      -float('inf'  )  )

        A_  :       Dict										=  1 - jnp.bool_(cur_len - self.max_length + 1  )

        A_  :       Dict										=  jnp.where(snake_case_					,      new_scores.at[:, self.eos_token_id].set(0  )					,      snake_case_  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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

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

        A_  :       Any										=  min_length
        A_  :       List[Any]										=  eos_token_id

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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

        A_  :       Optional[Any]										=  jnp.where(snake_case_					,      scores.at[:, self.eos_token_id].set(-float('inf'  )  )					,      snake_case_  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       List[Any]										=  list(snake_case_  )
        A_  :       Tuple										=  begin_index

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Dict										=  1 - jnp.bool_(cur_len - self.begin_index  )

        A_  :       int										=  jnp.where(snake_case_					,      scores.at[:, self.begin_suppress_tokens].set(-float('inf'  )  )					,      snake_case_  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       List[Any]										=  list(snake_case_  )

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  scores.at[..., self.suppress_tokens].set(-float('inf'  )  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Any										=  dict(snake_case_  )
        # 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.
        A_  :       Tuple										=  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:
            A_  :       Tuple										=  force_token_array.at[index].set(snake_case_  )
        A_  :       Any										=  jnp.intaa(snake_case_  )

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        def _force_token(snake_case_  ):
          A_  :       List[Any]										=  scores.shape[0]
          A_  :       Any										=  self.force_token_array[generation_idx]

          A_  :       Tuple										=  jnp.ones_like(snake_case_					,      dtype=scores.dtype  ) * -float('inf'  )
          A_  :       List[Any]										=  jnp.zeros((batch_size, 1)					,      dtype=scores.dtype  )
          A_  :       int										=  lax.dynamic_update_slice(snake_case_					,      snake_case_					,      (0, current_token)  )
          return new_scores

        A_  :       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(snake_case_  )					,      lambda: scores					,      )					,      )
        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Tuple										=  generate_config.eos_token_id
        A_  :       Optional[int]										=  generate_config.no_timestamps_token_id
        A_  :       List[str]										=  generate_config.no_timestamps_token_id + 1

        A_  :       Any										=  decoder_input_length + 1

        if generate_config.is_multilingual:
          # room for language token and task token
          self.begin_index += 2
        if hasattr(snake_case_					,      'max_initial_timestamp_index'  ):
          A_  :       List[Any]										=  generate_config.max_initial_timestamp_index
        else:
          A_  :       Any										=  model_config.vocab_size
        if self.max_initial_timestamp_index is None:
          A_  :       Optional[Any]										=  model_config.vocab_size

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       List[str]										=  scores.at[:, self.no_timestamps_token_id].set(-float('inf'  )  )

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

          A_  :       Tuple										=  jnp.where((cur_len - self.begin_index) < 2					,      snake_case_					,      snake_case_  )
          A_  :       Any										=  jnp.where(
              input_ids_k[cur_len - 2] >= self.timestamp_begin					,      snake_case_					,      snake_case_					,      )

          return jnp.where(
              snake_case_					,      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'  )  )					,      )					,      snake_case_					,      )

        A_  :       Tuple										=  jax.vmap(snake_case_  )(snake_case_					,      snake_case_  )

        A_  :       Optional[Any]										=  jnp.where(cur_len == self.begin_index					,      snake_case_					,      snake_case_  )
        A_  :       Tuple										=  jnp.where(
            self.max_initial_timestamp_index is not None					,      True and apply_max_initial_timestamp					,      snake_case_					,      )

        A_  :       int										=  self.timestamp_begin + self.max_initial_timestamp_index

        A_  :       List[Any]										=  jnp.where(
            snake_case_					,      scores.at[:, last_allowed + 1 :].set(-float('inf'  )  )					,      snake_case_					,      )

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

        def handle_cumulative_probs(snake_case_					,      snake_case_  ):
          A_  :       Dict										=  jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]					,      axis=-1  )
          A_  :       Optional[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'  )  )					,      snake_case_					,      )

        A_  :       Union[str, Any]										=  jax.vmap(snake_case_  )(snake_case_					,      snake_case_  )

        return scores