


'''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 A__       :






      def __init__( self			:Union[str, Any]     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:Optional[Any]=1_3     ,						SCREAMING_SNAKE_CASE			:str=7     ,						SCREAMING_SNAKE_CASE			:int=6     ,						SCREAMING_SNAKE_CASE			:Dict=1_7     ,						SCREAMING_SNAKE_CASE			:Dict=2_3     ,						SCREAMING_SNAKE_CASE			:List[Any]=1_1     ,						SCREAMING_SNAKE_CASE			:int=True     ,						)       ->					int:




             '''simple docstring'''
             _a    : Dict     					=parent
             _a    : str     					=batch_size
             _a    : List[Any]     					=seq_length
             _a    : Any     					=act_dim
             _a    : Dict     					=state_dim
             _a    : List[str]     					=hidden_size
             _a    : List[Any]     					=max_length
             _a    : Optional[Any]     					=is_training






      def __UpperCAmelCase       ( self			:List[Any]	)       ->					Tuple:




             '''simple docstring'''
             _a    : Union[str, Any]     					=floats_tensor((self.batch_size, self.seq_length, self.state_dim)	)
             _a    : Optional[Any]     					=floats_tensor((self.batch_size, self.seq_length, self.act_dim)	)
             _a    : Tuple     					=floats_tensor((self.batch_size, self.seq_length, 1)	)
             _a    : Optional[Any]     					=floats_tensor((self.batch_size, self.seq_length, 1)	)
             _a    : str     					=ids_tensor((self.batch_size, self.seq_length)     ,						vocab_size=1_0_0_0	)
             _a    : Optional[Any]     					=random_attention_mask((self.batch_size, self.seq_length)	)

             _a    : Dict     					=self.get_config()

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






      def __UpperCAmelCase       ( self			:List[str]	)       ->					Optional[Any]:




             '''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 __UpperCAmelCase       ( self			:Union[str, Any]     ,						SCREAMING_SNAKE_CASE			:List[Any]     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:List[Any]     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:str     ,						)       ->					List[str]:




             '''simple docstring'''
             _a    : List[str]     					=DecisionTransformerModel(config=_UpperCamelCase	)
             model.to(_UpperCamelCase	)
             model.eval()
             _a    : Dict     					=model(_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase	)

             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			:Any	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : List[Any]     					=self.prepare_config_and_inputs()
             (
                 _a  
             )    : List[str]     					=config_and_inputs
             _a    : 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 A__       (						_snake_case ,     _snake_case ,     _snake_case ,     unittest.TestCase      ):
      __UpperCamelCase       :  int	       =							(DecisionTransformerModel,) if is_torch_available() else ()
      __UpperCamelCase       :  Tuple	       =							()
      __UpperCamelCase       :  Any	       =							{'''feature-extraction''': DecisionTransformerModel} if is_torch_available() else {}

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

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






      def __UpperCAmelCase       ( self			:Tuple	)       ->					Tuple:




             '''simple docstring'''
             _a    : str     					=DecisionTransformerModelTester(self	)
             _a    : str     					=ConfigTester(self     ,						config_class=_UpperCamelCase     ,						hidden_size=3_7	)






      def __UpperCAmelCase       ( self			:Optional[Any]	)       ->					List[str]:




             '''simple docstring'''
             self.config_tester.run_common_tests()






      def __UpperCAmelCase       ( self			:List[str]	)       ->					Any:




             '''simple docstring'''
             _a    : Any     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*_UpperCamelCase	)






      @slow
      def __UpperCAmelCase       ( self			:int	)       ->					List[Any]:




             '''simple docstring'''
             for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a    : int     					=DecisionTransformerModel.from_pretrained(_UpperCamelCase	)
                    self.assertIsNotNone(_UpperCamelCase	)






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : Optional[Any]     					=self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    _a    : List[Any]     					=model_class(_UpperCamelCase	)
                    _a    : Optional[int]     					=inspect.signature(model.forward	)
                    # signature.parameters is an OrderedDict => so arg_names order is deterministic
                    _a    : Union[str, Any]     					=[*signature.parameters.keys()]

                    _a    : List[str]     					=[
                        'states',
                        'actions',
                        'rewards',
                        'returns_to_go',
                        'timesteps',
                        'attention_mask',
                    ]

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







@require_torch
class A__       (						unittest.TestCase      ):






      @slow
      def __UpperCAmelCase       ( self			:List[Any]	)       ->					int:




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

             _a    : Tuple     					=torch.tensor(
                 [[0.242_793, -0.28_693_074, 0.8_742_613], [0.67_815_274, -0.08_101_085, -0.12_952_147]]     ,						device=_UpperCamelCase	)

             _a    : Tuple     					=torch.tensor(_UpperCamelCase     ,						device=_UpperCamelCase     ,						dtype=torch.floataa	).reshape(1     ,						1     ,						1	)
             _a    : Optional[int]     					=state
             _a    : str     					=torch.zeros(1     ,						0     ,						config.act_dim     ,						device=_UpperCamelCase     ,						dtype=torch.floataa	)
             _a    : List[Any]     					=torch.zeros(1     ,						0     ,						device=_UpperCamelCase     ,						dtype=torch.floataa	)
             _a    : Optional[int]     					=torch.tensor(0     ,						device=_UpperCamelCase     ,						dtype=torch.long	).reshape(1     ,						1	)

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

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

                    with torch.no_grad():
                           _a    : Optional[int]     					=model(
                               states=_UpperCamelCase     ,						actions=_UpperCamelCase     ,						rewards=_UpperCamelCase     ,						returns_to_go=_UpperCamelCase     ,						timesteps=_UpperCamelCase     ,						attention_mask=_UpperCamelCase     ,						return_dict=_UpperCamelCase     ,						)

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

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










import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_convbert import ConvBertTokenizer


__UpperCAmelCase								=						logging.get_logger(__name__)

__UpperCAmelCase								=						{'vocab_file': 'vocab.txt'}

__UpperCAmelCase								=						{
    'vocab_file': {
        'YituTech/conv-bert-base': 'https://huggingface.co/YituTech/conv-bert-base/resolve/main/vocab.txt',
        'YituTech/conv-bert-medium-small': (
            'https://huggingface.co/YituTech/conv-bert-medium-small/resolve/main/vocab.txt'
        ),
        'YituTech/conv-bert-small': 'https://huggingface.co/YituTech/conv-bert-small/resolve/main/vocab.txt',
    }
}

__UpperCAmelCase								=						{
    'YituTech/conv-bert-base': 512,
    'YituTech/conv-bert-medium-small': 512,
    'YituTech/conv-bert-small': 512,
}


__UpperCAmelCase								=						{
    'YituTech/conv-bert-base': {'do_lower_case': True},
    'YituTech/conv-bert-medium-small': {'do_lower_case': True},
    'YituTech/conv-bert-small': {'do_lower_case': True},
}



class   lowerCamelCase      (_snake_case			):







		'''simple docstring'''



		_snake_case    :    Optional[int]       						=  VOCAB_FILES_NAMES
		_snake_case    :    int       						=  PRETRAINED_VOCAB_FILES_MAP
		_snake_case    :    Dict       						=  PRETRAINED_INIT_CONFIGURATION
		_snake_case    :    List[str]       						=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		_snake_case    :    Any       						=  ConvBertTokenizer



		def __init__(							self				,		_UpperCamelCase=None				,		_UpperCamelCase=None				,		_UpperCamelCase=True				,		_UpperCamelCase="[UNK]"				,		_UpperCamelCase="[SEP]"				,		_UpperCamelCase="[PAD]"				,		_UpperCamelCase="[CLS]"				,		_UpperCamelCase="[MASK]"				,		_UpperCamelCase=True				,		_UpperCamelCase=None				,		**_UpperCamelCase				,		)	->							Dict:
							super().__init__(
							    _UpperCamelCase				,		tokenizer_file=_UpperCamelCase				,		do_lower_case=_UpperCamelCase				,		unk_token=_UpperCamelCase				,		sep_token=_UpperCamelCase				,		pad_token=_UpperCamelCase				,		cls_token=_UpperCamelCase				,		mask_token=_UpperCamelCase				,		tokenize_chinese_chars=_UpperCamelCase				,		strip_accents=_UpperCamelCase				,		**_UpperCamelCase				,		)

							UpperCAmelCase_    :      Optional[int]            =						json.loads(self.backend_tokenizer.normalizer.__getstate__()   )
							if (
							    normalizer_state.get('lowercase'				,		_UpperCamelCase   ) != do_lower_case
							    or normalizer_state.get('strip_accents'				,		_UpperCamelCase   ) != strip_accents
							    or normalizer_state.get('handle_chinese_chars'				,		_UpperCamelCase   ) != tokenize_chinese_chars
							):
												UpperCAmelCase_    :      Any            =						getattr(_UpperCamelCase				,		normalizer_state.pop('type'   )   )
												UpperCAmelCase_    :      str            =						do_lower_case
												UpperCAmelCase_    :      List[Any]            =						strip_accents
												UpperCAmelCase_    :      str            =						tokenize_chinese_chars
												UpperCAmelCase_    :      Tuple            =						normalizer_class(**_UpperCamelCase   )

							UpperCAmelCase_    :      Any            =						do_lower_case



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase=None   )	->							List[str]:
							UpperCAmelCase_    :      int            =						[self.cls_token_id] + token_ids_a + [self.sep_token_id]

							if token_ids_a:
												output += token_ids_a + [self.sep_token_id]

							return output



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = None   )	->							List[int]:
							UpperCAmelCase_    :      Union[str, Any]            =						[self.sep_token_id]
							UpperCAmelCase_    :      int            =						[self.cls_token_id]
							if token_ids_a is None:
												return len(cls + token_ids_a + sep   ) * [0]
							return len(cls + token_ids_a + sep   ) * [0] + len(token_ids_a + sep   ) * [1]



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = None   )	->							Tuple[str]:
							UpperCAmelCase_    :      Any            =						self._tokenizer.model.save(_UpperCamelCase				,		name=_UpperCamelCase   )
							return tuple(_UpperCamelCase   )






