





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     __snake_case       :





       '''simple docstring'''







       def __init__(  self :    List[str]							,				A :    Dict							,				A :    Optional[Any]=13							,				A :    Optional[int]=7							,				A :    Optional[Any]=6							,				A :    List[Any]=17							,				A :    Any=23							,				A :    Union[str, Any]=11							,				A :    int=True							,				):
         __snake_case:						Tuple   =						parent
         __snake_case:						Optional[int]   =						batch_size
         __snake_case:						Optional[int]   =						seq_length
         __snake_case:						Optional[int]   =						act_dim
         __snake_case:						Tuple   =						state_dim
         __snake_case:						List[Any]   =						hidden_size
         __snake_case:						Tuple   =						max_length
         __snake_case:						List[str]   =						is_training




       def 							UpperCAmelCase__ (  self :    Dict		):
         __snake_case:						Optional[int]   =						floats_tensor((self.batch_size, self.seq_length, self.state_dim)		)
         __snake_case:						Any   =						floats_tensor((self.batch_size, self.seq_length, self.act_dim)		)
         __snake_case:						str   =						floats_tensor((self.batch_size, self.seq_length, 1)		)
         __snake_case:						Optional[Any]   =						floats_tensor((self.batch_size, self.seq_length, 1)		)
         __snake_case:						Dict   =						ids_tensor((self.batch_size, self.seq_length)							,				vocab_size=1_000		)
         __snake_case:						Optional[Any]   =						random_attention_mask((self.batch_size, self.seq_length)		)

         __snake_case:						Optional[Any]   =						self.get_config()

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




       def 							UpperCAmelCase__ (  self :    Optional[Any]		):
         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 							UpperCAmelCase__ (  self :    int							,				A :    Dict							,				A :    Optional[int]							,				A :    Any							,				A :    Any							,				A :    List[Any]							,				A :    Any							,				A :    Union[str, Any]							,				):
         __snake_case:						Tuple   =						DecisionTransformerModel(config=__a		)
         model.to(__a		)
         model.eval()
         __snake_case:						int   =						model(__a							,				__a							,				__a							,				__a							,				__a							,				__a		)

         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 							UpperCAmelCase__ (  self :    Tuple		):
         __snake_case:						Any   =						self.prepare_config_and_inputs()
         (
             (
             __snake_case  
         )  ,   (
             __snake_case  
         )  ,   (
             __snake_case  
         )  ,   (
             __snake_case  
         )  ,   (
             __snake_case  
         )  ,   (
             __snake_case  
         )  ,   (
             __snake_case  
         )  ,     
         ):						Tuple   =						config_and_inputs
         __snake_case:						Tuple   =						{
             """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     __snake_case       (							__lowerCamelCase   ,			__lowerCamelCase   ,			__lowerCamelCase   ,			unittest.TestCase			):





       '''simple docstring'''







       lowerCAmelCase__      		=     (DecisionTransformerModel,) if is_torch_available() else ()
       lowerCAmelCase__      		=     ()
       lowerCAmelCase__      		=     {"""feature-extraction""": DecisionTransformerModel} if is_torch_available() else {}

       # Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
       lowerCAmelCase__      		=     False

       # Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     False
       lowerCAmelCase__      		=     False




       def 							UpperCAmelCase__ (  self :    Optional[int]		):
         __snake_case:						Union[str, Any]   =						DecisionTransformerModelTester(self		)
         __snake_case:						Any   =						ConfigTester(self							,				config_class=__a							,				hidden_size=37		)




       def 							UpperCAmelCase__ (  self :    Union[str, Any]		):
         self.config_tester.run_common_tests()




       def 							UpperCAmelCase__ (  self :    Tuple		):
         __snake_case:						List[str]   =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__a		)




       @slow
       def 							UpperCAmelCase__ (  self :    Dict		):
         for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           __snake_case:						Dict   =						DecisionTransformerModel.from_pretrained(__a		)
           self.assertIsNotNone(__a		)






       def 							UpperCAmelCase__ (  self :    Tuple		):
         __snake_case  ,   __snake_case:						Optional[int]   =						self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           __snake_case:						Optional[int]   =						model_class(__a		)
           __snake_case:						Dict   =						inspect.signature(model.forward		)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           __snake_case:						Tuple   =						[*signature.parameters.keys()]

           __snake_case:						List[str]   =						[
               """states""",
               """actions""",
               """rewards""",
               """returns_to_go""",
               """timesteps""",
               """attention_mask""",
           ]

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







@require_torch
class     __snake_case       (							unittest.TestCase			):





       '''simple docstring'''







       @slow
       def 							UpperCAmelCase__ (  self :    Union[str, Any]		):
         __snake_case:						Tuple   =						2  # number of steps of autoregressive prediction we will perform
         __snake_case:						Tuple   =						10  # defined by the RL environment, may be normalized
         __snake_case:						Optional[int]   =						DecisionTransformerModel.from_pretrained("""edbeeching/decision-transformer-gym-hopper-expert"""		)
         __snake_case:						Optional[int]   =						model.to(__a		)
         __snake_case:						List[Any]   =						model.config
         torch.manual_seed(0		)
         __snake_case:						List[str]   =						torch.randn(1							,				1							,				config.state_dim		).to(device=__a							,				dtype=torch.floataa		)  # env.reset()

         __snake_case:						Union[str, Any]   =						torch.tensor(
             [[0.24_2793, -0.2869_3074, 0.874_2613], [0.6781_5274, -0.0810_1085, -0.1295_2147]]							,				device=__a		)

         __snake_case:						Dict   =						torch.tensor(__a							,				device=__a							,				dtype=torch.floataa		).reshape(1							,				1							,				1		)
         __snake_case:						List[str]   =						state
         __snake_case:						Optional[int]   =						torch.zeros(1							,				0							,				config.act_dim							,				device=__a							,				dtype=torch.floataa		)
         __snake_case:						List[str]   =						torch.zeros(1							,				0							,				device=__a							,				dtype=torch.floataa		)
         __snake_case:						Dict   =						torch.tensor(0							,				device=__a							,				dtype=torch.long		).reshape(1							,				1		)

         for step in range(__a		):
           __snake_case:						Optional[Any]   =						torch.cat([actions, torch.zeros(1							,				1							,				config.act_dim							,				device=__a		)]							,				dim=1		)
           __snake_case:						Dict   =						torch.cat([rewards, torch.zeros(1							,				1							,				device=__a		)]							,				dim=1		)

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

           with torch.no_grad():
             __snake_case  ,   __snake_case  ,   __snake_case:						Dict   =						model(
                 states=__a							,				actions=__a							,				rewards=__a							,				returns_to_go=__a							,				timesteps=__a							,				attention_mask=__a							,				return_dict=__a							,				)

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

           __snake_case:						str   =						action_pred[0, -1]
           __snake_case:						int   =						torch.cat([states, state]							,				dim=1		)
           __snake_case:						Optional[int]   =						returns_to_go[0, -1] - reward
           __snake_case:						Optional[Any]   =						torch.cat([returns_to_go, pred_return.reshape(1							,				1							,				1		)]							,				dim=1		)
           __snake_case:						List[str]   =						torch.cat(
               [timesteps, torch.ones((1, 1)							,				device=__a							,				dtype=torch.long		) * (step + 1)]							,				dim=1		)








"""simple docstring"""





import torch
from torch import nn






class     _UpperCAmelCase(    nn.Module	):



		def __init__(					self				,		__a				,		__a				,		__a				,		__a				,		__a=1				,		__a=False)       ->       Tuple:






							'''simple docstring'''
							super().__init__()

							_UpperCamelCase							 =  n_token
							_UpperCamelCase							 =  d_embed
							_UpperCamelCase							 =  d_proj

							_UpperCamelCase							 =  cutoffs + [n_token]
							_UpperCamelCase							 =  [0] + self.cutoffs
							_UpperCamelCase							 =  div_val

							_UpperCamelCase							 =  self.cutoffs[0]
							_UpperCamelCase							 =  len(self.cutoffs) - 1
							_UpperCamelCase							 =  self.shortlist_size + self.n_clusters

							if self.n_clusters > 0:
												_UpperCamelCase							 =  nn.Parameter(torch.zeros(self.n_clusters				,		self.d_embed))
												_UpperCamelCase							 =  nn.Parameter(torch.zeros(self.n_clusters))

							_UpperCamelCase							 =  nn.ModuleList()
							_UpperCamelCase							 =  nn.ParameterList()

							if div_val == 1:
												for i in range(len(self.cutoffs)):
																	if d_proj != d_embed:
																						self.out_projs.append(nn.Parameter(torch.FloatTensor(__a				,		__a)))
																	else:
																						self.out_projs.append(__a)

												self.out_layers.append(nn.Linear(__a				,		__a))
							else:
												for i in range(len(self.cutoffs)):
																	_UpperCamelCase ,  _UpperCamelCase							 =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
																	_UpperCamelCase							 =  d_embed // (div_val**i)

																	self.out_projs.append(nn.Parameter(torch.FloatTensor(__a				,		__a)))

																	self.out_layers.append(nn.Linear(__a				,		r_idx - l_idx))

							_UpperCamelCase							 =  keep_order



		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a)       ->       Any:






							'''simple docstring'''
							if proj is None:
												_UpperCamelCase							 =  nn.functional.linear(__a				,		__a				,		bias=__a)
							else:
												# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
												_UpperCamelCase							 =  nn.functional.linear(__a				,		proj.t().contiguous())
												_UpperCamelCase							 =  nn.functional.linear(__a				,		__a				,		bias=__a)
												# else:
												#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
												#     if bias is not None:
												#         logit = logit + bias

							return logit



		def 						UpperCAmelCase							(					self				,		__a				,		__a=None				,		__a=False)       ->       Dict:






							'''simple docstring'''

							if labels is not None:
												# Shift so that tokens < n predict n
												_UpperCamelCase							 =  hidden[..., :-1, :].contiguous()
												_UpperCamelCase							 =  labels[..., 1:].contiguous()
												_UpperCamelCase							 =  hidden.view(-1				,		hidden.size(-1))
												_UpperCamelCase							 =  labels.view(-1)
												if hidden.size(0) != labels.size(0):
																	raise RuntimeError('''Input and labels should have the same size in the batch dimension.''')
							else:
												_UpperCamelCase							 =  hidden.view(-1				,		hidden.size(-1))

							if self.n_clusters == 0:
												_UpperCamelCase							 =  self._compute_logit(__a				,		self.out_layers[0].weight				,		self.out_layers[0].bias				,		self.out_projs[0])
												if labels is not None:
																	_UpperCamelCase							 =  labels != -1_00
																	_UpperCamelCase							 =  torch.zeros_like(__a				,		dtype=hidden.dtype				,		device=hidden.device)
																	_UpperCamelCase							 =  (
																	    -nn.functional.log_softmax(__a				,		dim=-1)[mask].gather(1				,		labels[mask].unsqueeze(1)).squeeze(1)
																	)
												else:
																	_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=-1)
							else:
												# construct weights and biases
												_UpperCamelCase ,  _UpperCamelCase							 =  [], []
												for i in range(len(self.cutoffs)):
																	if self.div_val == 1:
																						_UpperCamelCase ,  _UpperCamelCase							 =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
																						_UpperCamelCase							 =  self.out_layers[0].weight[l_idx:r_idx]
																						_UpperCamelCase							 =  self.out_layers[0].bias[l_idx:r_idx]
																	else:
																						_UpperCamelCase							 =  self.out_layers[i].weight
																						_UpperCamelCase							 =  self.out_layers[i].bias

																	if i == 0:
																						_UpperCamelCase							 =  torch.cat([weight_i, self.cluster_weight]				,		dim=0)
																						_UpperCamelCase							 =  torch.cat([bias_i, self.cluster_bias]				,		dim=0)

																	weights.append(__a)
																	biases.append(__a)

												_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  weights[0], biases[0], self.out_projs[0]

												_UpperCamelCase							 =  self._compute_logit(__a				,		__a				,		__a				,		__a)
												_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=1)

												if labels is None:
																	_UpperCamelCase							 =  hidden.new_empty((head_logit.size(0), self.n_token))
												else:
																	_UpperCamelCase							 =  torch.zeros_like(__a				,		dtype=hidden.dtype				,		device=hidden.device)

												_UpperCamelCase							 =  0
												_UpperCamelCase							 =  [0] + self.cutoffs
												for i in range(len(__a) - 1):
																	_UpperCamelCase ,  _UpperCamelCase							 =  cutoff_values[i], cutoff_values[i + 1]

																	if labels is not None:
																						_UpperCamelCase							 =  (labels >= l_idx) & (labels < r_idx)
																						_UpperCamelCase							 =  mask_i.nonzero().squeeze()

																						if indices_i.numel() == 0:
																											continue

																						_UpperCamelCase							 =  labels.index_select(0				,		__a) - l_idx
																						_UpperCamelCase							 =  head_logprob.index_select(0				,		__a)
																						_UpperCamelCase							 =  hidden.index_select(0				,		__a)
																	else:
																						_UpperCamelCase							 =  hidden

																	if i == 0:
																						if labels is not None:
																											_UpperCamelCase							 =  head_logprob_i.gather(1				,		target_i[:, None]).squeeze(1)
																						else:
																											_UpperCamelCase							 =  head_logprob[:, : self.cutoffs[0]]
																	else:
																						_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  weights[i], biases[i], self.out_projs[i]

																						_UpperCamelCase							 =  self._compute_logit(__a				,		__a				,		__a				,		__a)
																						_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=1)
																						_UpperCamelCase							 =  self.cutoffs[0] + i - 1  # No probability for the head cluster
																						if labels is not None:
																											_UpperCamelCase							 =  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
																											    1				,		target_i[:, None]).squeeze(1)
																						else:
																											_UpperCamelCase							 =  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
																											_UpperCamelCase							 =  logprob_i

																	if labels is not None:
																						if (hasattr(self				,		'''keep_order''') and self.keep_order) or keep_order:
																											out.index_copy_(0				,		__a				,		-logprob_i)
																						else:
																											out[offset : offset + logprob_i.size(0)].copy_(-logprob_i)
																						offset += logprob_i.size(0)

							return out




		def 						UpperCAmelCase							(					self				,		__a)       ->       List[str]:






							'''simple docstring'''
							if self.n_clusters == 0:
												_UpperCamelCase							 =  self._compute_logit(__a				,		self.out_layers[0].weight				,		self.out_layers[0].bias				,		self.out_projs[0])
												return nn.functional.log_softmax(__a				,		dim=-1)
							else:
												# construct weights and biases
												_UpperCamelCase ,  _UpperCamelCase							 =  [], []
												for i in range(len(self.cutoffs)):
																	if self.div_val == 1:
																						_UpperCamelCase ,  _UpperCamelCase							 =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
																						_UpperCamelCase							 =  self.out_layers[0].weight[l_idx:r_idx]
																						_UpperCamelCase							 =  self.out_layers[0].bias[l_idx:r_idx]
																	else:
																						_UpperCamelCase							 =  self.out_layers[i].weight
																						_UpperCamelCase							 =  self.out_layers[i].bias

																	if i == 0:
																						_UpperCamelCase							 =  torch.cat([weight_i, self.cluster_weight]				,		dim=0)
																						_UpperCamelCase							 =  torch.cat([bias_i, self.cluster_bias]				,		dim=0)

																	weights.append(__a)
																	biases.append(__a)

												_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  weights[0], biases[0], self.out_projs[0]
												_UpperCamelCase							 =  self._compute_logit(__a				,		__a				,		__a				,		__a)

												_UpperCamelCase							 =  hidden.new_empty((head_logit.size(0), self.n_token))
												_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=1)

												_UpperCamelCase							 =  [0] + self.cutoffs
												for i in range(len(__a) - 1):
																	_UpperCamelCase ,  _UpperCamelCase							 =  cutoff_values[i], cutoff_values[i + 1]

																	if i == 0:
																						_UpperCamelCase							 =  head_logprob[:, : self.cutoffs[0]]
																	else:
																						_UpperCamelCase ,  _UpperCamelCase ,  _UpperCamelCase							 =  weights[i], biases[i], self.out_projs[i]

																						_UpperCamelCase							 =  self._compute_logit(__a				,		__a				,		__a				,		__a)
																						_UpperCamelCase							 =  nn.functional.log_softmax(__a				,		dim=1)

																						_UpperCamelCase							 =  head_logprob[:, -i] + tail_logprob_i
																						_UpperCamelCase							 =  logprob_i

												return out






