





import inspect
import unittest

from transformers import DecisionTransformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import DecisionTransformerModel
					from transformers.models.decision_transformer.modeling_decision_transformer import (
					    DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
					)

class     __SCREAMING_SNAKE_CASE   :






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase=13					,	__lowerCAmelCase=7					,	__lowerCAmelCase=6					,	__lowerCAmelCase=17					,	__lowerCAmelCase=23					,	__lowerCAmelCase=11					,	__lowerCAmelCase=True					,	):
														UpperCamelCase__			       =     parent
														UpperCamelCase__			       =     batch_size
														UpperCamelCase__			       =     seq_length
														UpperCamelCase__			       =     act_dim
														UpperCamelCase__			       =     state_dim
														UpperCamelCase__			       =     hidden_size
														UpperCamelCase__			       =     max_length
														UpperCamelCase__			       =     is_training






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     floats_tensor((self.batch_size, self.seq_length, self.state_dim)					)
														UpperCamelCase__			       =     floats_tensor((self.batch_size, self.seq_length, self.act_dim)					)
														UpperCamelCase__			       =     floats_tensor((self.batch_size, self.seq_length, 1)					)
														UpperCamelCase__			       =     floats_tensor((self.batch_size, self.seq_length, 1)					)
														UpperCamelCase__			       =     ids_tensor((self.batch_size, self.seq_length)					,	vocab_size=1000					)
														UpperCamelCase__			       =     random_attention_mask((self.batch_size, self.seq_length)					)

														UpperCamelCase__			       =     self.get_config()

														return (
														    config,
														    states,
														    actions,
														    rewards,
														    returns_to_go,
														    timesteps,
														    attention_mask,
														)






							def 			_lowerCamelCase						(  self					):
														return DecisionTransformerConfig(
														    batch_size=self.batch_size					,	seq_length=self.seq_length					,	act_dim=self.act_dim					,	state_dim=self.state_dim					,	hidden_size=self.hidden_size					,	max_length=self.max_length					,	)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	):
														UpperCamelCase__			       =     DecisionTransformerModel(config=__SCREAMING_SNAKE_CASE					)
														model.to(__SCREAMING_SNAKE_CASE					)
														model.eval()
														UpperCamelCase__			       =     model(__SCREAMING_SNAKE_CASE					,	__SCREAMING_SNAKE_CASE					,	__SCREAMING_SNAKE_CASE					,	__SCREAMING_SNAKE_CASE					,	__SCREAMING_SNAKE_CASE					,	__SCREAMING_SNAKE_CASE					)

														self.parent.assertEqual(result.state_preds.shape					,	states.shape					)
														self.parent.assertEqual(result.action_preds.shape					,	actions.shape					)
														self.parent.assertEqual(result.return_preds.shape					,	returns_to_go.shape					)
														self.parent.assertEqual(
														    result.last_hidden_state.shape					,	(self.batch_size, self.seq_length * 3, self.hidden_size)					)  # seq length *3 as there are 3 modelities: states, returns and actions







							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.prepare_config_and_inputs()
														(
														    (
														    UpperCamelCase__	
														)       ,       (
														    UpperCamelCase__	
														)       ,       (
														    UpperCamelCase__	
														)       ,       (
														    UpperCamelCase__	
														)       ,       (
														    UpperCamelCase__	
														)       ,       (
														    UpperCamelCase__	
														)       ,       (
														    UpperCamelCase__	
														)       ,       	
														)			       =     config_and_inputs
														UpperCamelCase__			       =     {
														    """states""": states,
														    """actions""": actions,
														    """rewards""": rewards,
														    """returns_to_go""": returns_to_go,
														    """timesteps""": timesteps,
														    """attention_mask""": attention_mask,
														}
														return config, inputs_dict






@require_torch
class     __SCREAMING_SNAKE_CASE   (       lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						unittest.TestCase    ):
							snake_case       : Any         =    (DecisionTransformerModel,) if is_torch_available() else ()
							snake_case       : Tuple         =    ()
							snake_case       : Dict         =    {"""feature-extraction""": DecisionTransformerModel} if is_torch_available() else {}

							# Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
							snake_case       : List[Any]         =    False

							# Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
							snake_case       : Dict         =    False
							snake_case       : str         =    False
							snake_case       : Optional[int]         =    False
							snake_case       : int         =    False
							snake_case       : Tuple         =    False
							snake_case       : List[str]         =    False
							snake_case       : Any         =    False
							snake_case       : Dict         =    False
							snake_case       : List[Any]         =    False






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     DecisionTransformerModelTester(self					)
														UpperCamelCase__			       =     ConfigTester(self					,	config_class=__SCREAMING_SNAKE_CASE					,	hidden_size=37					)






							def 			_lowerCamelCase						(  self					):
														self.config_tester.run_common_tests()






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE					)






							@slow
							def 			_lowerCamelCase						(  self					):

														for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					UpperCamelCase__			       =     DecisionTransformerModel.from_pretrained(__SCREAMING_SNAKE_CASE					)
																					self.assertIsNotNone(__SCREAMING_SNAKE_CASE					)







							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__       ,       UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					UpperCamelCase__			       =     model_class(__SCREAMING_SNAKE_CASE					)
																					UpperCamelCase__			       =     inspect.signature(model.forward					)
																					# signature.parameters is an OrderedDict => so arg_names order is deterministic
																					UpperCamelCase__			       =     [*signature.parameters.keys()]

																					UpperCamelCase__			       =     [
																					    """states""",
																					    """actions""",
																					    """rewards""",
																					    """returns_to_go""",
																					    """timesteps""",
																					    """attention_mask""",
																					]

																					self.assertListEqual(arg_names[: len(__SCREAMING_SNAKE_CASE					)]					,	__SCREAMING_SNAKE_CASE					)






@require_torch
class     __SCREAMING_SNAKE_CASE   (       unittest.TestCase    ):






							@slow
							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     2  # number of steps of autoregressive prediction we will perform
														UpperCamelCase__			       =     10  # defined by the RL environment, may be normalized
														UpperCamelCase__			       =     DecisionTransformerModel.from_pretrained("""edbeeching/decision-transformer-gym-hopper-expert"""					)
														UpperCamelCase__			       =     model.to(__SCREAMING_SNAKE_CASE					)
														UpperCamelCase__			       =     model.config
														torch.manual_seed(0					)
														UpperCamelCase__			       =     torch.randn(1					,	1					,	config.state_dim					).to(device=__SCREAMING_SNAKE_CASE					,	dtype=torch.floataa					)  # env.reset()

														UpperCamelCase__			       =     torch.tensor(
														    [[0.24_2793, -0.2869_3074, 0.874_2613], [0.6781_5274, -0.0810_1085, -0.1295_2147]]					,	device=__SCREAMING_SNAKE_CASE					)

														UpperCamelCase__			       =     torch.tensor(__SCREAMING_SNAKE_CASE					,	device=__SCREAMING_SNAKE_CASE					,	dtype=torch.floataa					).reshape(1					,	1					,	1					)
														UpperCamelCase__			       =     state
														UpperCamelCase__			       =     torch.zeros(1					,	0					,	config.act_dim					,	device=__SCREAMING_SNAKE_CASE					,	dtype=torch.floataa					)
														UpperCamelCase__			       =     torch.zeros(1					,	0					,	device=__SCREAMING_SNAKE_CASE					,	dtype=torch.floataa					)
														UpperCamelCase__			       =     torch.tensor(0					,	device=__SCREAMING_SNAKE_CASE					,	dtype=torch.long					).reshape(1					,	1					)

														for step in range(__SCREAMING_SNAKE_CASE					):
																					UpperCamelCase__			       =     torch.cat([actions, torch.zeros(1					,	1					,	config.act_dim					,	device=__SCREAMING_SNAKE_CASE					)]					,	dim=1					)
																					UpperCamelCase__			       =     torch.cat([rewards, torch.zeros(1					,	1					,	device=__SCREAMING_SNAKE_CASE					)]					,	dim=1					)

																					UpperCamelCase__			       =     torch.ones(1					,	states.shape[1]					).to(dtype=torch.long					,	device=states.device					)

																					with torch.no_grad():
																												UpperCamelCase__       ,       UpperCamelCase__       ,       UpperCamelCase__			       =     model(
																												    states=__SCREAMING_SNAKE_CASE					,	actions=__SCREAMING_SNAKE_CASE					,	rewards=__SCREAMING_SNAKE_CASE					,	returns_to_go=__SCREAMING_SNAKE_CASE					,	timesteps=__SCREAMING_SNAKE_CASE					,	attention_mask=__SCREAMING_SNAKE_CASE					,	return_dict=__SCREAMING_SNAKE_CASE					,	)

																					self.assertEqual(action_pred.shape					,	actions.shape					)
																					self.assertTrue(torch.allclose(action_pred[0, -1]					,	expected_outputs[step]					,	atol=1E-4					)					)
																					UpperCamelCase__       ,       UpperCamelCase__       ,       UpperCamelCase__       ,       UpperCamelCase__			       =     (  # env.step(action)
																					    torch.randn(1					,	1					,	config.state_dim					).to(device=__SCREAMING_SNAKE_CASE					,	dtype=torch.floataa					),
																					    1.0,
																					    False,
																					    {},
																					)

																					UpperCamelCase__			       =     action_pred[0, -1]
																					UpperCamelCase__			       =     torch.cat([states, state]					,	dim=1					)
																					UpperCamelCase__			       =     returns_to_go[0, -1] - reward
																					UpperCamelCase__			       =     torch.cat([returns_to_go, pred_return.reshape(1					,	1					,	1					)]					,	dim=1					)
																					UpperCamelCase__			       =     torch.cat(
																					    [timesteps, torch.ones((1, 1)					,	device=__SCREAMING_SNAKE_CASE					,	dtype=torch.long					) * (step + 1)]					,	dim=1					)











import warnings

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


UpperCamelCase__     =	logging.get_logger(__name__)

UpperCamelCase__     =	{
    "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/config.json",
    "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/config.json",
}

class     __SCREAMING_SNAKE_CASE   (       _a    ):
							snake_case       : Any         =    """xlnet"""
							snake_case       : Optional[Any]         =    ["""mems"""]
							snake_case       : Any         =    {
							    """n_token""": """vocab_size""",  # Backward compatibility
							    """hidden_size""": """d_model""",
							    """num_attention_heads""": """n_head""",
							    """num_hidden_layers""": """n_layer""",
							}






							def __init__(  self					,	__lowerCAmelCase=32000					,	__lowerCAmelCase=1024					,	__lowerCAmelCase=24					,	__lowerCAmelCase=16					,	__lowerCAmelCase=4096					,	__lowerCAmelCase="gelu"					,	__lowerCAmelCase=True					,	__lowerCAmelCase="bi"					,	__lowerCAmelCase=0.02					,	__lowerCAmelCase=1E-12					,	__lowerCAmelCase=0.1					,	__lowerCAmelCase=512					,	__lowerCAmelCase=None					,	__lowerCAmelCase=True					,	__lowerCAmelCase=False					,	__lowerCAmelCase=False					,	__lowerCAmelCase=-1					,	__lowerCAmelCase=False					,	__lowerCAmelCase="last"					,	__lowerCAmelCase=True					,	__lowerCAmelCase="tanh"					,	__lowerCAmelCase=0.1					,	__lowerCAmelCase=5					,	__lowerCAmelCase=5					,	__lowerCAmelCase=5					,	__lowerCAmelCase=1					,	__lowerCAmelCase=2					,	**__lowerCAmelCase					,	):
														UpperCamelCase__			       =     vocab_size
														UpperCamelCase__			       =     d_model
														UpperCamelCase__			       =     n_layer
														UpperCamelCase__			       =     n_head
														if d_model % n_head != 0:
																					raise ValueError(f"""'d_model % n_head' ({d_model % n_head}) should be equal to 0"""					)
														if "d_head" in kwargs:
																					if kwargs["d_head"] != d_model // n_head:
																												raise ValueError(
																												    f"""`d_head` ({kwargs["d_head"]}) should be equal to `d_model // n_head` ({d_model // n_head})"""					)
														UpperCamelCase__			       =     d_model // n_head
														UpperCamelCase__			       =     ff_activation
														UpperCamelCase__			       =     d_inner
														UpperCamelCase__			       =     untie_r
														UpperCamelCase__			       =     attn_type

														UpperCamelCase__			       =     initializer_range
														UpperCamelCase__			       =     layer_norm_eps

														UpperCamelCase__			       =     dropout
														UpperCamelCase__			       =     mem_len
														UpperCamelCase__			       =     reuse_len
														UpperCamelCase__			       =     bi_data
														UpperCamelCase__			       =     clamp_len
														UpperCamelCase__			       =     same_length

														UpperCamelCase__			       =     summary_type
														UpperCamelCase__			       =     summary_use_proj
														UpperCamelCase__			       =     summary_activation
														UpperCamelCase__			       =     summary_last_dropout
														UpperCamelCase__			       =     start_n_top
														UpperCamelCase__			       =     end_n_top

														UpperCamelCase__			       =     bos_token_id
														UpperCamelCase__			       =     pad_token_id
														UpperCamelCase__			       =     eos_token_id

														if "use_cache" in kwargs:
																					warnings.warn(
																					    """The `use_cache` argument is deprecated and will be removed in a future version, use `use_mems_eval`"""
																					    """ instead."""					,	__lowerCAmelCase					,	)
																					UpperCamelCase__			       =     kwargs["""use_cache"""]

														UpperCamelCase__			       =     use_mems_eval
														UpperCamelCase__			       =     use_mems_train
														super().__init__(pad_token_id=__lowerCAmelCase					,	bos_token_id=__lowerCAmelCase					,	eos_token_id=__lowerCAmelCase					,	**__lowerCAmelCase					)






							@property
							def 			_lowerCamelCase						(  self					):
														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 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														# 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."""					)




