


'''simple docstring'''






import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
  import jax
  import jax.numpy as jnp

  from transformers.generation import (
      FlaxForcedBOSTokenLogitsProcessor,
      FlaxForcedEOSTokenLogitsProcessor,
      FlaxLogitsProcessorList,
      FlaxMinLengthLogitsProcessor,
      FlaxTemperatureLogitsWarper,
      FlaxTopKLogitsWarper,
      FlaxTopPLogitsWarper,
  )






@require_flax
class    __snake_case       (   unittest.TestCase):







    """simple docstring"""



    def __lowercase   (   self    :  List[str]    ,						lowerCamelCase    :  int    ,						lowerCamelCase    :  int		)   ->  Union[str, Any]:
      lowerCAmelCase_			:      Any			=   jnp.ones((batch_size, length)		) / length
      return scores

    def __lowercase   (   self    :  Dict		)   ->  Optional[int]:
      lowerCAmelCase_			:      Union[str, Any]			=   None
      lowerCAmelCase_			:      Any			=   20

      lowerCAmelCase_			:      str			=   self._get_uniform_logits(batch_size=2    ,						length=lowerCamelCase		)

      # tweak scores to not be uniform anymore
      lowerCAmelCase_			:      Optional[int]			=   scores.at[1, 5].set((1 / length) + 0.1		)  # peak, 1st batch
      lowerCAmelCase_			:      int			=   scores.at[1, 10].set((1 / length) - 0.4		)  # valley, 1st batch

      # compute softmax
      lowerCAmelCase_			:      List[str]			=   jax.nn.softmax(lowerCamelCase    ,						axis=-1		)

      lowerCAmelCase_			:      Optional[int]			=   FlaxTemperatureLogitsWarper(temperature=0.5		)
      lowerCAmelCase_			:      str			=   FlaxTemperatureLogitsWarper(temperature=1.3		)

      lowerCAmelCase_			:      Any			=   jax.nn.softmax(temp_dist_warper_sharper(lowerCamelCase    ,						scores.copy()    ,						cur_len=lowerCamelCase		)    ,						axis=-1		)
      lowerCAmelCase_			:      int			=   jax.nn.softmax(temp_dist_warper_smoother(lowerCamelCase    ,						scores.copy()    ,						cur_len=lowerCamelCase		)    ,						axis=-1		)

      # uniform distribution stays uniform
      self.assertTrue(jnp.allclose(probs[0, :]    ,						warped_prob_sharp[0, :]    ,						atol=1E-3		)		)
      self.assertTrue(jnp.allclose(probs[0, :]    ,						warped_prob_smooth[0, :]    ,						atol=1E-3		)		)

      # sharp peaks get higher, valleys get lower
      self.assertLess(probs[1, :].max()    ,						warped_prob_sharp[1, :].max()		)
      self.assertGreater(probs[1, :].min()    ,						warped_prob_sharp[1, :].min()		)

      # smooth peaks get lower, valleys get higher
      self.assertGreater(probs[1, :].max()    ,						warped_prob_smooth[1, :].max()		)
      self.assertLess(probs[1, :].min()    ,						warped_prob_smooth[1, :].min()		)

    def __lowercase   (   self    :  Optional[Any]		)   ->  Optional[int]:
      lowerCAmelCase_			:      Optional[Any]			=   None
      lowerCAmelCase_			:      List[Any]			=   10
      lowerCAmelCase_			:      List[str]			=   2

      # create ramp distribution
      lowerCAmelCase_			:      Union[str, Any]			=   np.broadcast_to(np.arange(lowerCamelCase		)[None, :]    ,						(batch_size, vocab_size)		).copy()
      lowerCAmelCase_			:      str			=   ramp_logits[1:, : vocab_size // 2] + vocab_size

      lowerCAmelCase_			:      Optional[int]			=   FlaxTopKLogitsWarper(3		)

      lowerCAmelCase_			:      int			=   top_k_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)

      # check that correct tokens are filtered
      self.assertListEqual(jnp.isinf(scores[0]		).tolist()    ,						7 * [True] + 3 * [False]		)
      self.assertListEqual(jnp.isinf(scores[1]		).tolist()    ,						2 * [True] + 3 * [False] + 5 * [True]		)

      # check special case
      lowerCAmelCase_			:      Optional[Any]			=   5
      lowerCAmelCase_			:      int			=   FlaxTopKLogitsWarper(top_k=1    ,						filter_value=0.0    ,						min_tokens_to_keep=3		)

      lowerCAmelCase_			:      Tuple			=   np.broadcast_to(np.arange(lowerCamelCase		)[None, :]    ,						(batch_size, length)		).copy()
      lowerCAmelCase_			:      Dict			=   top_k_warp_safety_check(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)

      # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
      self.assertListEqual((scores == 0.0).sum(axis=-1		).tolist()    ,						[2, 2]		)

    def __lowercase   (   self    :  Optional[int]		)   ->  Any:
      lowerCAmelCase_			:      Any			=   None
      lowerCAmelCase_			:      Dict			=   10
      lowerCAmelCase_			:      str			=   2

      # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
      lowerCAmelCase_			:      Dict			=   np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]		)		)

      lowerCAmelCase_			:      str			=   FlaxTopPLogitsWarper(0.8		)
      lowerCAmelCase_			:      str			=   np.exp(top_p_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)		)

      # dist should be filtered to keep min num values so that sum is >= top_p
      # exp (-inf) => 0
      lowerCAmelCase_			:      Optional[Any]			=   np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]		)
      self.assertTrue(np.allclose(lowerCamelCase    ,						lowerCamelCase    ,						atol=1E-3		)		)

      # check edge cases with negative and extreme logits
      lowerCAmelCase_			:      Dict			=   np.broadcast_to(np.arange(lowerCamelCase		)[None, :]    ,						(batch_size, vocab_size)		).copy() - (
          vocab_size // 2
      )

      # make ramp_logits more extreme
      lowerCAmelCase_			:      Optional[int]			=   ramp_logits[1] * 100.0

      # make sure at least 2 tokens are kept
      lowerCAmelCase_			:      Tuple			=   FlaxTopPLogitsWarper(0.9    ,						min_tokens_to_keep=2    ,						filter_value=0.0		)
      lowerCAmelCase_			:      Optional[Any]			=   top_p_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)

      # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
      self.assertListEqual((filtered_dist != 0.0).sum(axis=-1		).tolist()    ,						[3, 2]		)

    def __lowercase   (   self    :  Optional[Any]		)   ->  Tuple:
      lowerCAmelCase_			:      Optional[Any]			=   20
      lowerCAmelCase_			:      List[str]			=   4
      lowerCAmelCase_			:      Any			=   0

      lowerCAmelCase_			:      Dict			=   FlaxMinLengthLogitsProcessor(min_length=10    ,						eos_token_id=lowerCamelCase		)

      # check that min length is applied at length 5
      lowerCAmelCase_			:      Optional[int]			=   ids_tensor((batch_size, 20)    ,						vocab_size=20		)
      lowerCAmelCase_			:      Any			=   5
      lowerCAmelCase_			:      Tuple			=   self._get_uniform_logits(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Optional[Any]			=   min_dist_processor(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()    ,						4 * [-float("""inf"""		)]		)

      # check that min length is not applied anymore at length 15
      lowerCAmelCase_			:      str			=   self._get_uniform_logits(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Union[str, Any]			=   15
      lowerCAmelCase_			:      Optional[int]			=   min_dist_processor(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      self.assertFalse(jnp.isinf(lowerCamelCase		).any()		)

    def __lowercase   (   self    :  List[str]		)   ->  List[Any]:
      lowerCAmelCase_			:      List[str]			=   20
      lowerCAmelCase_			:      List[str]			=   4
      lowerCAmelCase_			:      str			=   0

      lowerCAmelCase_			:      List[Any]			=   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase		)

      # check that all scores are -inf except the bos_token_id score
      lowerCAmelCase_			:      List[str]			=   ids_tensor((batch_size, 1)    ,						vocab_size=20		)
      lowerCAmelCase_			:      int			=   1
      lowerCAmelCase_			:      int			=   self._get_uniform_logits(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      int			=   logits_processor(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]		).all()		)
      self.assertListEqual(scores[:, bos_token_id].tolist()    ,						4 * [0]		)  # score for bos_token_id shold be zero

      # check that bos_token_id is not forced if current length is greater than 1
      lowerCAmelCase_			:      List[str]			=   3
      lowerCAmelCase_			:      Tuple			=   self._get_uniform_logits(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      List[str]			=   logits_processor(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      self.assertFalse(jnp.isinf(lowerCamelCase		).any()		)

    def __lowercase   (   self    :  Union[str, Any]		)   ->  int:
      lowerCAmelCase_			:      str			=   20
      lowerCAmelCase_			:      List[str]			=   4
      lowerCAmelCase_			:      int			=   0
      lowerCAmelCase_			:      Any			=   5

      lowerCAmelCase_			:      Optional[Any]			=   FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase    ,						eos_token_id=lowerCamelCase		)

      # check that all scores are -inf except the eos_token_id when max_length is reached
      lowerCAmelCase_			:      Union[str, Any]			=   ids_tensor((batch_size, 4)    ,						vocab_size=20		)
      lowerCAmelCase_			:      Any			=   4
      lowerCAmelCase_			:      str			=   self._get_uniform_logits(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Optional[Any]			=   logits_processor(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]		).all()		)
      self.assertListEqual(scores[:, eos_token_id].tolist()    ,						4 * [0]		)  # score for eos_token_id should be zero

      # check that eos_token_id is not forced if max_length is not reached
      lowerCAmelCase_			:      List[str]			=   3
      lowerCAmelCase_			:      str			=   self._get_uniform_logits(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      int			=   logits_processor(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      self.assertFalse(jnp.isinf(lowerCamelCase		).any()		)

    def __lowercase   (   self    :  str		)   ->  List[Any]:
      lowerCAmelCase_			:      Optional[int]			=   4
      lowerCAmelCase_			:      Optional[Any]			=   10
      lowerCAmelCase_			:      Any			=   15
      lowerCAmelCase_			:      str			=   2
      lowerCAmelCase_			:      Dict			=   1
      lowerCAmelCase_			:      Tuple			=   15

      # dummy input_ids and scores
      lowerCAmelCase_			:      List[Any]			=   ids_tensor((batch_size, sequence_length)    ,						lowerCamelCase		)
      lowerCAmelCase_			:      List[str]			=   input_ids.copy()

      lowerCAmelCase_			:      int			=   self._get_uniform_logits(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Any			=   scores.copy()

      # instantiate all dist processors
      lowerCAmelCase_			:      Union[str, Any]			=   FlaxTemperatureLogitsWarper(temperature=0.5		)
      lowerCAmelCase_			:      Union[str, Any]			=   FlaxTopKLogitsWarper(3		)
      lowerCAmelCase_			:      List[str]			=   FlaxTopPLogitsWarper(0.8		)

      # instantiate all logits processors
      lowerCAmelCase_			:      List[Any]			=   FlaxMinLengthLogitsProcessor(min_length=10    ,						eos_token_id=lowerCamelCase		)
      lowerCAmelCase_			:      Dict			=   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase		)
      lowerCAmelCase_			:      Optional[Any]			=   FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase    ,						eos_token_id=lowerCamelCase		)

      lowerCAmelCase_			:      List[str]			=   10

      # no processor list
      lowerCAmelCase_			:      Dict			=   temp_dist_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      lowerCAmelCase_			:      Union[str, Any]			=   top_k_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   top_p_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      lowerCAmelCase_			:      Tuple			=   min_dist_proc(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      lowerCAmelCase_			:      Any			=   bos_dist_proc(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
      lowerCAmelCase_			:      Dict			=   eos_dist_proc(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)

      # with processor list
      lowerCAmelCase_			:      Optional[Any]			=   FlaxLogitsProcessorList(
          [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
      lowerCAmelCase_			:      List[Any]			=   processor(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)

      # scores should be equal
      self.assertTrue(jnp.allclose(lowerCamelCase    ,						lowerCamelCase    ,						atol=1E-3		)		)

      # input_ids should never be changed
      self.assertListEqual(input_ids.tolist()    ,						input_ids_comp.tolist()		)





    def __lowercase   (   self    :  int		)   ->  Optional[int]:
      lowerCAmelCase_			:      int			=   4
      lowerCAmelCase_			:      List[Any]			=   10
      lowerCAmelCase_			:      Union[str, Any]			=   15
      lowerCAmelCase_			:      Tuple			=   2
      lowerCAmelCase_			:      Optional[Any]			=   1
      lowerCAmelCase_			:      Union[str, Any]			=   15

      # dummy input_ids and scores
      lowerCAmelCase_			:      Tuple			=   ids_tensor((batch_size, sequence_length)    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Any			=   input_ids.copy()

      lowerCAmelCase_			:      List[Any]			=   self._get_uniform_logits(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      str			=   scores.copy()

      # instantiate all dist processors
      lowerCAmelCase_			:      Optional[int]			=   FlaxTemperatureLogitsWarper(temperature=0.5		)
      lowerCAmelCase_			:      Union[str, Any]			=   FlaxTopKLogitsWarper(3		)
      lowerCAmelCase_			:      int			=   FlaxTopPLogitsWarper(0.8		)

      # instantiate all logits processors
      lowerCAmelCase_			:      Tuple			=   FlaxMinLengthLogitsProcessor(min_length=10    ,						eos_token_id=lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase		)
      lowerCAmelCase_			:      str			=   FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase    ,						eos_token_id=lowerCamelCase		)

      lowerCAmelCase_			:      List[Any]			=   10

      # no processor list
      def run_no_processor_list(lowerCamelCase    :  int    ,						lowerCamelCase    :  Dict    ,						lowerCamelCase    :  Tuple		):
        lowerCAmelCase_			:      int			=   temp_dist_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
        lowerCAmelCase_			:      str			=   top_k_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
        lowerCAmelCase_			:      List[Any]			=   top_p_warp(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
        lowerCAmelCase_			:      Any			=   min_dist_proc(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
        lowerCAmelCase_			:      Tuple			=   bos_dist_proc(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
        lowerCAmelCase_			:      Tuple			=   eos_dist_proc(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
        return scores

      # with processor list
      def run_processor_list(lowerCamelCase    :  Dict    ,						lowerCamelCase    :  str    ,						lowerCamelCase    :  List[Any]		):
        lowerCAmelCase_			:      str			=   FlaxLogitsProcessorList(
            [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
        lowerCAmelCase_			:      Optional[int]			=   processor(lowerCamelCase    ,						lowerCamelCase    ,						cur_len=lowerCamelCase		)
        return scores

      lowerCAmelCase_			:      Optional[int]			=   jax.jit(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   jax.jit(lowerCamelCase		)

      lowerCAmelCase_			:      Any			=   jitted_run_no_processor_list(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      List[Any]			=   jitted_run_processor_list(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase		)

      # scores should be equal
      self.assertTrue(jnp.allclose(lowerCamelCase    ,						lowerCamelCase    ,						atol=1E-3		)		)

      # input_ids should never be changed
      self.assertListEqual(input_ids.tolist()    ,						input_ids_comp.tolist()		)







'''simple docstring'''






import json
import os
import unittest

from transformers.models.biogpt.tokenization_biogpt import VOCAB_FILES_NAMES, BioGptTokenizer
from transformers.testing_utils import slow

from ...test_tokenization_common import TokenizerTesterMixin






class    __snake_case       (   _SCREAMING_SNAKE_CASE	,unittest.TestCase):







    """simple docstring"""



    lowercase    =					BioGptTokenizer
    lowercase    =					False

    def __lowercase   (   self    :  List[Any]		)   ->  List[str]:
      super().setUp()

      # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
      lowerCAmelCase_			:      Optional[Any]			=   [
          """l""",
          """o""",
          """w""",
          """e""",
          """r""",
          """s""",
          """t""",
          """i""",
          """d""",
          """n""",
          """w</w>""",
          """r</w>""",
          """t</w>""",
          """lo""",
          """low""",
          """er</w>""",
          """low</w>""",
          """lowest</w>""",
          """newer</w>""",
          """wider</w>""",
          """<unk>""",
      ]
      lowerCAmelCase_			:      int			=   dict(zip(lowerCamelCase    ,						range(len(lowerCamelCase		)		)		)		)
      lowerCAmelCase_			:      int			=   ["""l o 123""", """lo w 1456""", """e r</w> 1789""", """"""]

      lowerCAmelCase_			:      Optional[int]			=   os.path.join(self.tmpdirname    ,						VOCAB_FILES_NAMES["""vocab_file"""]		)
      lowerCAmelCase_			:      Union[str, Any]			=   os.path.join(self.tmpdirname    ,						VOCAB_FILES_NAMES["""merges_file"""]		)
      with open(self.vocab_file    ,						"""w"""		) as fp:
        fp.write(json.dumps(lowerCamelCase		)		)
      with open(self.merges_file    ,						"""w"""		) as fp:
        fp.write("""\n""".join(lowerCamelCase		)		)

    def __lowercase   (   self    :  Union[str, Any]    ,						lowerCamelCase    :  Any		)   ->  List[str]:
      lowerCAmelCase_			:      Dict			=   """lower newer"""
      lowerCAmelCase_			:      str			=   """lower newer"""
      return input_text, output_text

    def __lowercase   (   self    :  Optional[int]		)   ->  str:
      lowerCAmelCase_			:      Any			=   BioGptTokenizer(self.vocab_file    ,						self.merges_file		)

      lowerCAmelCase_			:      int			=   """lower"""
      lowerCAmelCase_			:      str			=   ["""low""", """er</w>"""]
      lowerCAmelCase_			:      Any			=   tokenizer.tokenize(lowerCamelCase		)
      self.assertListEqual(lowerCamelCase    ,						lowerCamelCase		)

      lowerCAmelCase_			:      Dict			=   tokens + ["""<unk>"""]
      lowerCAmelCase_			:      Optional[Any]			=   [14, 15, 20]
      self.assertListEqual(tokenizer.convert_tokens_to_ids(lowerCamelCase		)    ,						lowerCamelCase		)





    @slow
    def __lowercase   (   self    :  str		)   ->  Optional[Any]:
      lowerCAmelCase_			:      Dict			=   BioGptTokenizer.from_pretrained("""microsoft/biogpt"""		)

      lowerCAmelCase_			:      int			=   tokenizer.encode("""sequence builders"""    ,						add_special_tokens=lowerCamelCase		)
      lowerCAmelCase_			:      Optional[Any]			=   tokenizer.encode("""multi-sequence build"""    ,						add_special_tokens=lowerCamelCase		)

      lowerCAmelCase_			:      str			=   tokenizer.build_inputs_with_special_tokens(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   tokenizer.build_inputs_with_special_tokens(lowerCamelCase    ,						lowerCamelCase		)

      self.assertTrue(encoded_sentence == [2] + text		)
      self.assertTrue(encoded_pair == [2] + text + [2] + text_a		)



