





"""simple docstring"""





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 					_lowerCAmelCase   :
       """simple docstring"""




       def __init__(    self    , __UpperCAmelCase    , __UpperCAmelCase=1_3    , __UpperCAmelCase=7    , __UpperCAmelCase=6    , __UpperCAmelCase=1_7    , __UpperCAmelCase=2_3    , __UpperCAmelCase=1_1    , __UpperCAmelCase=True    , ):



         '''simple docstring'''

         lowerCAmelCase__			:Dict											=							parent
         lowerCAmelCase__			:Optional[Any]											=							batch_size
         lowerCAmelCase__			:int											=							seq_length
         lowerCAmelCase__			:str											=							act_dim
         lowerCAmelCase__			:Optional[Any]											=							state_dim
         lowerCAmelCase__			:List[str]											=							hidden_size
         lowerCAmelCase__			:str											=							max_length
         lowerCAmelCase__			:Optional[int]											=							is_training






       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:Dict											=							floats_tensor((self.batch_size, self.seq_length, self.state_dim)      )
         lowerCAmelCase__			:List[Any]											=							floats_tensor((self.batch_size, self.seq_length, self.act_dim)      )
         lowerCAmelCase__			:Union[str, Any]											=							floats_tensor((self.batch_size, self.seq_length, 1)      )
         lowerCAmelCase__			:List[str]											=							floats_tensor((self.batch_size, self.seq_length, 1)      )
         lowerCAmelCase__			:List[Any]											=							ids_tensor((self.batch_size, self.seq_length)    , vocab_size=1_0_0_0      )
         lowerCAmelCase__			:List[str]											=							random_attention_mask((self.batch_size, self.seq_length)      )

         lowerCAmelCase__			:Tuple											=							self.get_config()

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






       def      snake_case						(    self      ):



         '''simple docstring'''

         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      snake_case						(    self    , __UpperCAmelCase    , __UpperCAmelCase    , __UpperCAmelCase    , __UpperCAmelCase    , __UpperCAmelCase    , __UpperCAmelCase    , __UpperCAmelCase    , ):



         '''simple docstring'''

         lowerCAmelCase__			:Union[str, Any]											=							DecisionTransformerModel(config=__a      )
         model.to(__a      )
         model.eval()
         lowerCAmelCase__			: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      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:List[Any]											=							self.prepare_config_and_inputs()
         (
             (
             lowerCAmelCase__					
         )							,				(
             lowerCAmelCase__					
         )							,				(
             lowerCAmelCase__					
         )							,				(
             lowerCAmelCase__					
         )							,				(
             lowerCAmelCase__					
         )							,				(
             lowerCAmelCase__					
         )							,				(
             lowerCAmelCase__					
         )							,									
         )			:Any											=							config_and_inputs
         lowerCAmelCase__			:Optional[Any]											=							{
             '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 					_lowerCAmelCase   (     lowerCamelCase_  ,   lowerCamelCase_  ,   lowerCamelCase_  ,   unittest.TestCase       ):
       """simple docstring"""




       __magic_name__     :List[str]		  =					(DecisionTransformerModel,) if is_torch_available() else ()
       __magic_name__     :Tuple		  =					()
       __magic_name__     :Tuple		  =					{"""feature-extraction""": DecisionTransformerModel} if is_torch_available() else {}

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

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






       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:Union[str, Any]											=							DecisionTransformerModelTester(self      )
         lowerCAmelCase__			:Optional[int]											=							ConfigTester(self    , config_class=__a    , hidden_size=3_7      )






       def      snake_case						(    self      ):



         '''simple docstring'''

         self.config_tester.run_common_tests()






       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:Dict											=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__a      )






       @slow
       def      snake_case						(    self      ):



         '''simple docstring'''

         for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowerCAmelCase__			:List[Any]											=							DecisionTransformerModel.from_pretrained(__a      )
           self.assertIsNotNone(__a      )






       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__							,				lowerCAmelCase__			:Optional[int]											=							self.model_tester.prepare_config_and_inputs_for_common()

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

           lowerCAmelCase__			:Tuple											=							[
               'states',
               'actions',
               'rewards',
               'returns_to_go',
               'timesteps',
               'attention_mask',
           ]

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






@require_torch
class 					_lowerCAmelCase   (     unittest.TestCase       ):
       """simple docstring"""




       @slow
       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:int											=							2  # number of steps of autoregressive prediction we will perform
         lowerCAmelCase__			:Any											=							1_0  # defined by the RL environment, may be normalized
         lowerCAmelCase__			:List[str]											=							DecisionTransformerModel.from_pretrained('edbeeching/decision-transformer-gym-hopper-expert'      )
         lowerCAmelCase__			:List[str]											=							model.to(__a      )
         lowerCAmelCase__			:List[Any]											=							model.config
         torch.manual_seed(0      )
         lowerCAmelCase__			:str											=							torch.randn(1    , 1    , config.state_dim      ).to(device=__a    , dtype=torch.floataa      )  # env.reset()

         lowerCAmelCase__			:Any											=							torch.tensor(
             [[0.24_27_93, -0.28_69_30_74, 0.8_74_26_13], [0.67_81_52_74, -0.08_10_10_85, -0.12_95_21_47]]    , device=__a      )

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

         for step in range(__a      ):
           lowerCAmelCase__			:Tuple											=							torch.cat([actions, torch.zeros(1    , 1    , config.act_dim    , device=__a      )]    , dim=1      )
           lowerCAmelCase__			:Union[str, Any]											=							torch.cat([rewards, torch.zeros(1    , 1    , device=__a      )]    , dim=1      )

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

           with torch.no_grad():
             lowerCAmelCase__							,				lowerCAmelCase__							,				lowerCAmelCase__			:str											=							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      )      )
           lowerCAmelCase__							,				lowerCAmelCase__							,				lowerCAmelCase__							,				lowerCAmelCase__			:Union[str, Any]											=							(  # env.step(action)
               torch.randn(1    , 1    , config.state_dim      ).to(device=__a    , dtype=torch.floataa      ),
               1.0,
               False,
               {},
           )

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







'''simple docstring'''







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


lowerCAmelCase_					:    Dict      	=			logging.get_logger(__name__)

lowerCAmelCase_					:    int      	=			{
    'bigcode/gpt_bigcode-santacoder': 'https://huggingface.co/bigcode/gpt_bigcode-santacoder/resolve/main/config.json',
}
class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_      ):






				"""simple docstring"""







				__a            ='gpt_bigcode'
				__a            =['past_key_values']
				__a            ={
				    'hidden_size': 'n_embd',
				    'max_position_embeddings': 'n_positions',
				    'num_attention_heads': 'n_head',
				    'num_hidden_layers': 'n_layer',
				}




				def __init__(  self     :					Optional[Any]					,				__a     :					Tuple=5_02_57					,				__a     :					str=10_24					,				__a     :					Dict=7_68					,				__a     :					Tuple=12					,				__a     :					str=12					,				__a     :					Optional[int]=None					,				__a     :					Dict="gelu_pytorch_tanh"					,				__a     :					Tuple=0.1					,				__a     :					Tuple=0.1					,				__a     :					Union[str, Any]=0.1					,				__a     :					Tuple=1e-5					,				__a     :					str=0.02					,				__a     :					Dict=True					,				__a     :					Union[str, Any]=True					,				__a     :					Optional[int]=5_02_56					,				__a     :					Optional[int]=5_02_56					,				__a     :					Union[str, Any]=True					,				__a     :					Dict=True					,				__a     :					Union[str, Any]=True					,				**__a     :					List[Any]					,				):
								_a								= vocab_size
								_a								= n_positions
								_a								= n_embd
								_a								= n_layer
								_a								= n_head
								_a								= n_inner
								_a								= activation_function
								_a								= resid_pdrop
								_a								= embd_pdrop
								_a								= attn_pdrop
								_a								= layer_norm_epsilon
								_a								= initializer_range
								_a								= scale_attn_weights
								_a								= use_cache
								_a								= attention_softmax_in_fpaa
								_a								= scale_attention_softmax_in_fpaa
								_a								= multi_query

								_a								= bos_token_id
								_a								= eos_token_id

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