

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 	UpperCAmelCase__		(	unittest.TestCase	):






  """simple docstring"""


  def 				lowercase_		(							self						:		List[Any]     ,       __lowerCamelCase						:		int     ,       __lowerCamelCase						:		int )		->							Optional[Any]:
       SCREAMING_SNAKE_CASE__   =   jnp.ones((batch_size, length) ) / length
       return scores






  def 				lowercase_		(							self						:		Any )		->							Tuple:
       SCREAMING_SNAKE_CASE__   =   None
       SCREAMING_SNAKE_CASE__   =   20

       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(batch_size=2     ,       length=__lowerCamelCase )

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

       # compute softmax
       SCREAMING_SNAKE_CASE__   =   jax.nn.softmax(__lowerCamelCase     ,       axis=-1 )

       SCREAMING_SNAKE_CASE__   =   FlaxTemperatureLogitsWarper(temperature=0.5 )
       SCREAMING_SNAKE_CASE__   =   FlaxTemperatureLogitsWarper(temperature=1.3 )

       SCREAMING_SNAKE_CASE__   =   jax.nn.softmax(temp_dist_warper_sharper(__lowerCamelCase     ,       scores.copy()     ,       cur_len=__lowerCamelCase )     ,       axis=-1 )
       SCREAMING_SNAKE_CASE__   =   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						:		List[str] )		->							Tuple:
       SCREAMING_SNAKE_CASE__   =   None
       SCREAMING_SNAKE_CASE__   =   10
       SCREAMING_SNAKE_CASE__   =   2

       # create ramp distribution
       SCREAMING_SNAKE_CASE__   =   np.broadcast_to(np.arange(__lowerCamelCase )[None, :]     ,       (batch_size, vocab_size) ).copy()
       SCREAMING_SNAKE_CASE__   =   ramp_logits[1:, : vocab_size // 2] + vocab_size

       SCREAMING_SNAKE_CASE__   =   FlaxTopKLogitsWarper(3 )

       SCREAMING_SNAKE_CASE__   =   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
       SCREAMING_SNAKE_CASE__   =   5
       SCREAMING_SNAKE_CASE__   =   FlaxTopKLogitsWarper(top_k=1     ,       filter_value=0.0     ,       min_tokens_to_keep=3 )

       SCREAMING_SNAKE_CASE__   =   np.broadcast_to(np.arange(__lowerCamelCase )[None, :]     ,       (batch_size, length) ).copy()
       SCREAMING_SNAKE_CASE__   =   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						:		List[str] )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   None
       SCREAMING_SNAKE_CASE__   =   10
       SCREAMING_SNAKE_CASE__   =   2

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

       SCREAMING_SNAKE_CASE__   =   FlaxTopPLogitsWarper(0.8 )
       SCREAMING_SNAKE_CASE__   =   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
       SCREAMING_SNAKE_CASE__   =   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
       SCREAMING_SNAKE_CASE__   =   np.broadcast_to(np.arange(__lowerCamelCase )[None, :]     ,       (batch_size, vocab_size) ).copy() - (
           vocab_size // 2
       )

       # make ramp_logits more extreme
       SCREAMING_SNAKE_CASE__   =   ramp_logits[1] * 100.0

       # make sure at least 2 tokens are kept
       SCREAMING_SNAKE_CASE__   =   FlaxTopPLogitsWarper(0.9     ,       min_tokens_to_keep=2     ,       filter_value=0.0 )
       SCREAMING_SNAKE_CASE__   =   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						:		str )		->							List[Any]:
       SCREAMING_SNAKE_CASE__   =   20
       SCREAMING_SNAKE_CASE__   =   4
       SCREAMING_SNAKE_CASE__   =   0

       SCREAMING_SNAKE_CASE__   =   FlaxMinLengthLogitsProcessor(min_length=10     ,       eos_token_id=__lowerCamelCase )

       # check that min length is applied at length 5
       SCREAMING_SNAKE_CASE__   =   ids_tensor((batch_size, 20)     ,       vocab_size=20 )
       SCREAMING_SNAKE_CASE__   =   5
       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   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
       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   15
       SCREAMING_SNAKE_CASE__   =   min_dist_processor(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
       self.assertFalse(jnp.isinf(__lowerCamelCase ).any() )






  def 				lowercase_		(							self						:		Union[str, Any] )		->							Dict:
       SCREAMING_SNAKE_CASE__   =   20
       SCREAMING_SNAKE_CASE__   =   4
       SCREAMING_SNAKE_CASE__   =   0

       SCREAMING_SNAKE_CASE__   =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__lowerCamelCase )

       # check that all scores are -inf except the bos_token_id score
       SCREAMING_SNAKE_CASE__   =   ids_tensor((batch_size, 1)     ,       vocab_size=20 )
       SCREAMING_SNAKE_CASE__   =   1
       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   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
       SCREAMING_SNAKE_CASE__   =   3
       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   logits_processor(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
       self.assertFalse(jnp.isinf(__lowerCamelCase ).any() )






  def 				lowercase_		(							self						:		Union[str, Any] )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   20
       SCREAMING_SNAKE_CASE__   =   4
       SCREAMING_SNAKE_CASE__   =   0
       SCREAMING_SNAKE_CASE__   =   5

       SCREAMING_SNAKE_CASE__   =   FlaxForcedEOSTokenLogitsProcessor(max_length=__lowerCamelCase     ,       eos_token_id=__lowerCamelCase )

       # check that all scores are -inf except the eos_token_id when max_length is reached
       SCREAMING_SNAKE_CASE__   =   ids_tensor((batch_size, 4)     ,       vocab_size=20 )
       SCREAMING_SNAKE_CASE__   =   4
       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   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
       SCREAMING_SNAKE_CASE__   =   3
       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   logits_processor(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
       self.assertFalse(jnp.isinf(__lowerCamelCase ).any() )






  def 				lowercase_		(							self						:		str )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   4
       SCREAMING_SNAKE_CASE__   =   10
       SCREAMING_SNAKE_CASE__   =   15
       SCREAMING_SNAKE_CASE__   =   2
       SCREAMING_SNAKE_CASE__   =   1
       SCREAMING_SNAKE_CASE__   =   15

       # dummy input_ids and scores
       SCREAMING_SNAKE_CASE__   =   ids_tensor((batch_size, sequence_length)     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   input_ids.copy()

       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   scores.copy()

       # instantiate all dist processors
       SCREAMING_SNAKE_CASE__   =   FlaxTemperatureLogitsWarper(temperature=0.5 )
       SCREAMING_SNAKE_CASE__   =   FlaxTopKLogitsWarper(3 )
       SCREAMING_SNAKE_CASE__   =   FlaxTopPLogitsWarper(0.8 )

       # instantiate all logits processors
       SCREAMING_SNAKE_CASE__   =   FlaxMinLengthLogitsProcessor(min_length=10     ,       eos_token_id=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   FlaxForcedEOSTokenLogitsProcessor(max_length=__lowerCamelCase     ,       eos_token_id=__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   10

       # no processor list
       SCREAMING_SNAKE_CASE__   =   temp_dist_warp(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   top_k_warp(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   top_p_warp(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   min_dist_proc(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   bos_dist_proc(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   eos_dist_proc(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )

       # with processor list
       SCREAMING_SNAKE_CASE__   =   FlaxLogitsProcessorList(
           [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
       SCREAMING_SNAKE_CASE__   =   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						:		List[str] )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   4
       SCREAMING_SNAKE_CASE__   =   10
       SCREAMING_SNAKE_CASE__   =   15
       SCREAMING_SNAKE_CASE__   =   2
       SCREAMING_SNAKE_CASE__   =   1
       SCREAMING_SNAKE_CASE__   =   15

       # dummy input_ids and scores
       SCREAMING_SNAKE_CASE__   =   ids_tensor((batch_size, sequence_length)     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   input_ids.copy()

       SCREAMING_SNAKE_CASE__   =   self._get_uniform_logits(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   scores.copy()

       # instantiate all dist processors
       SCREAMING_SNAKE_CASE__   =   FlaxTemperatureLogitsWarper(temperature=0.5 )
       SCREAMING_SNAKE_CASE__   =   FlaxTopKLogitsWarper(3 )
       SCREAMING_SNAKE_CASE__   =   FlaxTopPLogitsWarper(0.8 )

       # instantiate all logits processors
       SCREAMING_SNAKE_CASE__   =   FlaxMinLengthLogitsProcessor(min_length=10     ,       eos_token_id=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   FlaxForcedEOSTokenLogitsProcessor(max_length=__lowerCamelCase     ,       eos_token_id=__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   10

       # no processor list
       def run_no_processor_list(__lowerCamelCase						:		Tuple     ,       __lowerCamelCase						:		Tuple     ,       __lowerCamelCase						:		List[Any] ):
            SCREAMING_SNAKE_CASE__   =   temp_dist_warp(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   top_k_warp(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   top_p_warp(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   min_dist_proc(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   bos_dist_proc(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   eos_dist_proc(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
            return scores

       # with processor list
       def run_processor_list(__lowerCamelCase						:		Union[str, Any]     ,       __lowerCamelCase						:		Any     ,       __lowerCamelCase						:		List[str] ):
            SCREAMING_SNAKE_CASE__   =   FlaxLogitsProcessorList(
                [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
            SCREAMING_SNAKE_CASE__   =   processor(__lowerCamelCase     ,       __lowerCamelCase     ,       cur_len=__lowerCamelCase )
            return scores

       SCREAMING_SNAKE_CASE__   =   jax.jit(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   jax.jit(__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   jitted_run_no_processor_list(__lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   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() )





from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE					:  List[str]             =      logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE					:  Tuple             =      {
    '''transfo-xl-wt103''': '''https://huggingface.co/transfo-xl-wt103/resolve/main/config.json''',
}






class 	UpperCAmelCase__		(	A__	):






  """simple docstring"""


  a									= "transfo-xl"
  a									= ["mems"]
  a									= {
      "n_token": "vocab_size",
      "hidden_size": "d_model",
      "num_attention_heads": "n_head",
      "num_hidden_layers": "n_layer",
  }






  def __init__(							self						:		Any     ,       __lowerCamelCase						:		int=26_7735     ,       __lowerCamelCase						:		Any=[2_0000, 4_0000, 20_0000]     ,       __lowerCamelCase						:		Dict=1024     ,       __lowerCamelCase						:		Optional[int]=1024     ,       __lowerCamelCase						:		List[Any]=16     ,       __lowerCamelCase						:		Union[str, Any]=64     ,       __lowerCamelCase						:		Dict=4096     ,       __lowerCamelCase						:		int=4     ,       __lowerCamelCase						:		Dict=False     ,       __lowerCamelCase						:		Tuple=18     ,       __lowerCamelCase						:		Optional[int]=1600     ,       __lowerCamelCase						:		str=1000     ,       __lowerCamelCase						:		Any=True     ,       __lowerCamelCase						:		Dict=True     ,       __lowerCamelCase						:		List[Any]=0     ,       __lowerCamelCase						:		int=-1     ,       __lowerCamelCase						:		Tuple=True     ,       __lowerCamelCase						:		List[Any]=0.1     ,       __lowerCamelCase						:		int=0.0     ,       __lowerCamelCase						:		int=True     ,       __lowerCamelCase						:		str="normal"     ,       __lowerCamelCase						:		List[str]=0.01     ,       __lowerCamelCase						:		Any=0.01     ,       __lowerCamelCase						:		List[str]=0.02     ,       __lowerCamelCase						:		List[str]=1e-5     ,       __lowerCamelCase						:		Union[str, Any]=0     ,       **__lowerCamelCase						:		int     ,       )		->							Optional[Any]:
       SCREAMING_SNAKE_CASE__   =   vocab_size
       SCREAMING_SNAKE_CASE__   =   []
       self.cutoffs.extend(__lowerCamelCase )
       if proj_share_all_but_first:
            SCREAMING_SNAKE_CASE__   =   [False] + [True] * len(self.cutoffs )
       else:
            SCREAMING_SNAKE_CASE__   =   [False] + [False] * len(self.cutoffs )
       SCREAMING_SNAKE_CASE__   =   d_model
       SCREAMING_SNAKE_CASE__   =   d_embed
       SCREAMING_SNAKE_CASE__   =   d_head
       SCREAMING_SNAKE_CASE__   =   d_inner
       SCREAMING_SNAKE_CASE__   =   div_val
       SCREAMING_SNAKE_CASE__   =   pre_lnorm
       SCREAMING_SNAKE_CASE__   =   n_layer
       SCREAMING_SNAKE_CASE__   =   n_head
       SCREAMING_SNAKE_CASE__   =   mem_len
       SCREAMING_SNAKE_CASE__   =   same_length
       SCREAMING_SNAKE_CASE__   =   attn_type
       SCREAMING_SNAKE_CASE__   =   clamp_len
       SCREAMING_SNAKE_CASE__   =   sample_softmax
       SCREAMING_SNAKE_CASE__   =   adaptive
       SCREAMING_SNAKE_CASE__   =   dropout
       SCREAMING_SNAKE_CASE__   =   dropatt
       SCREAMING_SNAKE_CASE__   =   untie_r
       SCREAMING_SNAKE_CASE__   =   init
       SCREAMING_SNAKE_CASE__   =   init_range
       SCREAMING_SNAKE_CASE__   =   proj_init_std
       SCREAMING_SNAKE_CASE__   =   init_std
       SCREAMING_SNAKE_CASE__   =   layer_norm_epsilon
       super().__init__(eos_token_id=__lowerCamelCase     ,       **__lowerCamelCase )






  @property
  def 				lowercase_		(							self						:		str )		->							Dict:
       # Message copied from Transformer-XL documentation
       logger.info(f'''The model {self.model_type} is one of the few models that has no sequence length limit.''' )
       return -1







  @max_position_embeddings.setter
  def 				lowercase_		(							self						:		List[str]     ,       __lowerCamelCase						:		Any )		->							List[Any]:
       # Message copied from Transformer-XL documentation
       raise NotImplementedError(
           f'''The model {self.model_type} is one of the few models that has no sequence length limit.''' )


