



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







    def 			lowercase_       (      self :  Optional[Any]						,			__snake_case :  str						,			__snake_case :  Dict	):
      a						:					Optional[Any]        =     jnp.ones((batch_size, length)	) / length
      return scores







    def 			lowercase_       (      self :  str	):
      a						:					Tuple        =     None
      a						:					Optional[int]        =     20

      a						:					List[str]        =     self._get_uniform_logits(batch_size=2						,			length=_A	)

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

      # compute softmax
      a						:					str        =     jax.nn.softmax(_A						,			axis=-1	)

      a						:					List[str]        =     FlaxTemperatureLogitsWarper(temperature=0.5	)
      a						:					List[Any]        =     FlaxTemperatureLogitsWarper(temperature=1.3	)

      a						:					Tuple        =     jax.nn.softmax(temp_dist_warper_sharper(_A						,			scores.copy()						,			cur_len=_A	)						,			axis=-1	)
      a						:					str        =     jax.nn.softmax(temp_dist_warper_smoother(_A						,			scores.copy()						,			cur_len=_A	)						,			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 :  Any	):
      a						:					int        =     None
      a						:					Optional[int]        =     10
      a						:					List[str]        =     2

      # create ramp distribution
      a						:					int        =     np.broadcast_to(np.arange(_A	)[None, :]						,			(batch_size, vocab_size)	).copy()
      a						:					List[Any]        =     ramp_logits[1:, : vocab_size // 2] + vocab_size

      a						:					List[Any]        =     FlaxTopKLogitsWarper(3	)

      a						:					str        =     top_k_warp(_A						,			_A						,			cur_len=_A	)

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

      a						:					str        =     np.broadcast_to(np.arange(_A	)[None, :]						,			(batch_size, length)	).copy()
      a						:					Dict        =     top_k_warp_safety_check(_A						,			_A						,			cur_len=_A	)

      # 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[Any]	):
      a						:					Optional[Any]        =     None
      a						:					Optional[Any]        =     10
      a						:					str        =     2

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

      a						:					Optional[int]        =     FlaxTopPLogitsWarper(0.8	)
      a						:					Dict        =     np.exp(top_p_warp(_A						,			_A						,			cur_len=_A	)	)

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

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

      # make ramp_logits more extreme
      a						:					str        =     ramp_logits[1] * 1_00.0

      # make sure at least 2 tokens are kept
      a						:					List[str]        =     FlaxTopPLogitsWarper(0.9						,			min_tokens_to_keep=2						,			filter_value=0.0	)
      a						:					Union[str, Any]        =     top_p_warp(_A						,			_A						,			cur_len=_A	)

      # 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 :  int	):
      a						:					Dict        =     20
      a						:					Union[str, Any]        =     4
      a						:					Optional[int]        =     0

      a						:					List[Any]        =     FlaxMinLengthLogitsProcessor(min_length=10						,			eos_token_id=_A	)

      # check that min length is applied at length 5
      a						:					Dict        =     ids_tensor((batch_size, 20)						,			vocab_size=20	)
      a						:					str        =     5
      a						:					List[Any]        =     self._get_uniform_logits(_A						,			_A	)
      a						:					List[Any]        =     min_dist_processor(_A						,			_A						,			cur_len=_A	)
      self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()						,			4 * [-float('inf'	)]	)

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







    def 			lowercase_       (      self :  Any	):
      a						:					List[str]        =     20
      a						:					Dict        =     4
      a						:					str        =     0

      a						:					List[str]        =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A	)

      # check that all scores are -inf except the bos_token_id score
      a						:					Tuple        =     ids_tensor((batch_size, 1)						,			vocab_size=20	)
      a						:					List[str]        =     1
      a						:					List[str]        =     self._get_uniform_logits(_A						,			_A	)
      a						:					Dict        =     logits_processor(_A						,			_A						,			cur_len=_A	)
      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
      a						:					int        =     3
      a						:					Any        =     self._get_uniform_logits(_A						,			_A	)
      a						:					List[str]        =     logits_processor(_A						,			_A						,			cur_len=_A	)
      self.assertFalse(jnp.isinf(_A	).any()	)







    def 			lowercase_       (      self :  Dict	):
      a						:					List[Any]        =     20
      a						:					str        =     4
      a						:					Optional[int]        =     0
      a						:					Optional[Any]        =     5

      a						:					List[Any]        =     FlaxForcedEOSTokenLogitsProcessor(max_length=_A						,			eos_token_id=_A	)

      # check that all scores are -inf except the eos_token_id when max_length is reached
      a						:					str        =     ids_tensor((batch_size, 4)						,			vocab_size=20	)
      a						:					int        =     4
      a						:					Dict        =     self._get_uniform_logits(_A						,			_A	)
      a						:					List[str]        =     logits_processor(_A						,			_A						,			cur_len=_A	)
      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
      a						:					str        =     3
      a						:					Dict        =     self._get_uniform_logits(_A						,			_A	)
      a						:					Tuple        =     logits_processor(_A						,			_A						,			cur_len=_A	)
      self.assertFalse(jnp.isinf(_A	).any()	)







    def 			lowercase_       (      self :  Any	):
      a						:					str        =     4
      a						:					int        =     10
      a						:					str        =     15
      a						:					Optional[int]        =     2
      a						:					Optional[Any]        =     1
      a						:					List[Any]        =     15

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

      a						:					Union[str, Any]        =     self._get_uniform_logits(_A						,			_A	)
      a						:					List[Any]        =     scores.copy()

      # instantiate all dist processors
      a						:					str        =     FlaxTemperatureLogitsWarper(temperature=0.5	)
      a						:					Any        =     FlaxTopKLogitsWarper(3	)
      a						:					str        =     FlaxTopPLogitsWarper(0.8	)

      # instantiate all logits processors
      a						:					Union[str, Any]        =     FlaxMinLengthLogitsProcessor(min_length=10						,			eos_token_id=_A	)
      a						:					int        =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A	)
      a						:					Any        =     FlaxForcedEOSTokenLogitsProcessor(max_length=_A						,			eos_token_id=_A	)

      a						:					str        =     10

      # no processor list
      a						:					Tuple        =     temp_dist_warp(_A						,			_A						,			cur_len=_A	)
      a						:					Any        =     top_k_warp(_A						,			_A						,			cur_len=_A	)
      a						:					List[Any]        =     top_p_warp(_A						,			_A						,			cur_len=_A	)
      a						:					Dict        =     min_dist_proc(_A						,			_A						,			cur_len=_A	)
      a						:					Optional[Any]        =     bos_dist_proc(_A						,			_A						,			cur_len=_A	)
      a						:					Any        =     eos_dist_proc(_A						,			_A						,			cur_len=_A	)

      # with processor list
      a						:					str        =     FlaxLogitsProcessorList(
          [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]	)
      a						:					Any        =     processor(_A						,			_A						,			cur_len=_A	)

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

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







    def 			lowercase_       (      self :  Union[str, Any]	):
      a						:					Tuple        =     4
      a						:					Optional[Any]        =     10
      a						:					Tuple        =     15
      a						:					str        =     2
      a						:					Tuple        =     1
      a						:					str        =     15

      # dummy input_ids and scores
      a						:					int        =     ids_tensor((batch_size, sequence_length)						,			_A	)
      a						:					List[str]        =     input_ids.copy()

      a						:					Union[str, Any]        =     self._get_uniform_logits(_A						,			_A	)
      a						:					Optional[Any]        =     scores.copy()

      # instantiate all dist processors
      a						:					List[Any]        =     FlaxTemperatureLogitsWarper(temperature=0.5	)
      a						:					Optional[int]        =     FlaxTopKLogitsWarper(3	)
      a						:					int        =     FlaxTopPLogitsWarper(0.8	)

      # instantiate all logits processors
      a						:					List[Any]        =     FlaxMinLengthLogitsProcessor(min_length=10						,			eos_token_id=_A	)
      a						:					Any        =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A	)
      a						:					int        =     FlaxForcedEOSTokenLogitsProcessor(max_length=_A						,			eos_token_id=_A	)

      a						:					List[Any]        =     10

      # no processor list
      def run_no_processor_list(__snake_case :  int						,			__snake_case :  Optional[int]						,			__snake_case :  List[Any]	):
        a						:					Optional[Any]        =     temp_dist_warp(_A						,			_A						,			cur_len=_A	)
        a						:					Optional[Any]        =     top_k_warp(_A						,			_A						,			cur_len=_A	)
        a						:					int        =     top_p_warp(_A						,			_A						,			cur_len=_A	)
        a						:					Tuple        =     min_dist_proc(_A						,			_A						,			cur_len=_A	)
        a						:					List[Any]        =     bos_dist_proc(_A						,			_A						,			cur_len=_A	)
        a						:					Optional[int]        =     eos_dist_proc(_A						,			_A						,			cur_len=_A	)
        return scores

      # with processor list
      def run_processor_list(__snake_case :  Union[str, Any]						,			__snake_case :  Union[str, Any]						,			__snake_case :  Any	):
        a						:					Union[str, Any]        =     FlaxLogitsProcessorList(
            [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]	)
        a						:					Optional[Any]        =     processor(_A						,			_A						,			cur_len=_A	)
        return scores

      a						:					Optional[Any]        =     jax.jit(_A	)
      a						:					Dict        =     jax.jit(_A	)

      a						:					List[Any]        =     jitted_run_no_processor_list(_A						,			_A						,			_A	)
      a						:					int        =     jitted_run_processor_list(_A						,			_A						,			_A	)

      # scores should be equal
      self.assertTrue(jnp.allclose(_A						,			_A						,			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


UpperCamelCase__    			=			logging.get_logger(__name__)

UpperCamelCase__    			=			{
    """edbeeching/decision-transformer-gym-hopper-medium""": (
        """https://huggingface.co/edbeeching/decision-transformer-gym-hopper-medium/resolve/main/config.json"""
    ),
    # See all DecisionTransformer models at https://huggingface.co/models?filter=decision_transformer
}
class 			a__    (					snake_case__      ):
						_a							:    Optional[int]					=    """decision_transformer"""
						_a							:    Optional[int]					=    ["""past_key_values"""]
						_a							:    Dict					=    {
						    """max_position_embeddings""": """n_positions""",
						    """num_attention_heads""": """n_head""",
						    """num_hidden_layers""": """n_layer""",
						}






						def __init__(   self		,			_A=1_7		,			_A=4		,			_A=1_2_8		,			_A=4_0_9_6		,			_A=True		,			_A=1		,			_A=1_0_2_4		,			_A=3		,			_A=1		,			_A=None		,			_A="relu"		,			_A=0.1		,			_A=0.1		,			_A=0.1		,			_A=1E-5		,			_A=0.02		,			_A=True		,			_A=True		,			_A=5_0_2_5_6		,			_A=5_0_2_5_6		,			_A=False		,			_A=False		,			**_A		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       state_dim
											__lowerCAmelCase				   =       act_dim
											__lowerCAmelCase				   =       hidden_size
											__lowerCAmelCase				   =       max_ep_len
											__lowerCAmelCase				   =       action_tanh
											__lowerCAmelCase				   =       vocab_size
											__lowerCAmelCase				   =       n_positions
											__lowerCAmelCase				   =       n_layer
											__lowerCAmelCase				   =       n_head
											__lowerCAmelCase				   =       n_inner
											__lowerCAmelCase				   =       activation_function
											__lowerCAmelCase				   =       resid_pdrop
											__lowerCAmelCase				   =       embd_pdrop
											__lowerCAmelCase				   =       attn_pdrop
											__lowerCAmelCase				   =       layer_norm_epsilon
											__lowerCAmelCase				   =       initializer_range
											__lowerCAmelCase				   =       scale_attn_weights
											__lowerCAmelCase				   =       use_cache
											__lowerCAmelCase				   =       scale_attn_by_inverse_layer_idx
											__lowerCAmelCase				   =       reorder_and_upcast_attn

											__lowerCAmelCase				   =       bos_token_id
											__lowerCAmelCase				   =       eos_token_id

											super().__init__(bos_token_id=_A		,			eos_token_id=_A		,			**_A			)




