
"""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						:					Union[str, Any]       ,   _snake_case						:					List[str]       ,   _snake_case						:					Union[str, Any]=13       ,   _snake_case						:					str=7       ,   _snake_case						:					List[Any]=6       ,   _snake_case						:					Optional[Any]=17       ,   _snake_case						:					Tuple=23       ,   _snake_case						:					int=11       ,   _snake_case						:					Optional[int]=True       ,   )      -> List[str]:







								"""simple docstring"""
								A_				 =						parent
								A_				 =						batch_size
								A_				 =						seq_length
								A_				 =						act_dim
								A_				 =						state_dim
								A_				 =						hidden_size
								A_				 =						max_length
								A_				 =						is_training




			def      lowerCamelCase__     (				self						:					Union[str, Any]	)      -> Union[str, Any]:







								"""simple docstring"""
								A_				 =						floats_tensor((self.batch_size, self.seq_length, self.state_dim)	)
								A_				 =						floats_tensor((self.batch_size, self.seq_length, self.act_dim)	)
								A_				 =						floats_tensor((self.batch_size, self.seq_length, 1)	)
								A_				 =						floats_tensor((self.batch_size, self.seq_length, 1)	)
								A_				 =						ids_tensor((self.batch_size, self.seq_length)       ,   vocab_size=1_000	)
								A_				 =						random_attention_mask((self.batch_size, self.seq_length)	)

								A_				 =						self.get_config()

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




			def      lowerCamelCase__     (				self						:					int	)      -> Union[str, 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      lowerCamelCase__     (				self						:					int       ,   _snake_case						:					List[str]       ,   _snake_case						:					Union[str, Any]       ,   _snake_case						:					Dict       ,   _snake_case						:					Any       ,   _snake_case						:					List[str]       ,   _snake_case						:					Any       ,   _snake_case						:					str       ,   )      -> Dict:







								"""simple docstring"""
								A_				 =						DecisionTransformerModel(config=__A	)
								model.to(__A	)
								model.eval()
								A_				 =						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      lowerCamelCase__     (				self						:					List[Any]	)      -> str:







								"""simple docstring"""
								A_				 =						self.prepare_config_and_inputs()
								(
								    A_		
								)				 =						config_and_inputs
								A_				 =						{
								    "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				(					_lowercase     ,				_lowercase     ,				_lowercase     ,				unittest.TestCase      ):




			"""simple docstring"""







			snake_case	      =				(DecisionTransformerModel,) if is_torch_available() else ()
			snake_case	      =				()
			snake_case	      =				{"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	      =				False

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




			def      lowerCamelCase__     (				self						:					int	)      -> List[Any]:







								"""simple docstring"""
								A_				 =						DecisionTransformerModelTester(self	)
								A_				 =						ConfigTester(self       ,   config_class=__A       ,   hidden_size=37	)




			def      lowerCamelCase__     (				self						:					Tuple	)      -> str:







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




			def      lowerCamelCase__     (				self						:					Optional[Any]	)      -> Optional[Any]:







								"""simple docstring"""
								A_				 =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__A	)




			@slow
			def      lowerCamelCase__     (				self						:					Optional[Any]	)      -> Tuple:







								"""simple docstring"""
								for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													A_				 =						DecisionTransformerModel.from_pretrained(__A	)
													self.assertIsNotNone(__A	)






			def      lowerCamelCase__     (				self						:					Union[str, Any]	)      -> Union[str, Any]:







								"""simple docstring"""
								A_				 =						self.model_tester.prepare_config_and_inputs_for_common()

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

													A_				 =						[
													    "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      lowerCamelCase__     (				self						:					str	)      -> List[Any]:







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

								A_				 =						torch.tensor(
								    [[0.2_4_2_7_9_3, -0.2_8_6_9_3_0_7_4, 0.8_7_4_2_6_1_3], [0.6_7_8_1_5_2_7_4, -0.0_8_1_0_1_0_8_5, -0.1_2_9_5_2_1_4_7]]       ,   device=__A	)

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

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

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

													with torch.no_grad():
																		A_				 =						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	)	)
													A_				 =						(  # env.step(action)
													    torch.randn(1       ,   1       ,   config.state_dim	).to(device=__A       ,   dtype=torch.floataa	),
													    1.0,
													    False,
													    {},
													)

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






"""simple docstring"""





import inspect
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_config_docstrings.py
UpperCamelCase_					:						Optional[Any]				= '''src/transformers'''


# This is to make sure the transformers module imported is the one in the repo.
UpperCamelCase_					:						Union[str, Any]				= direct_transformers_import(PATH_TO_TRANSFORMERS)

UpperCamelCase_					:						List[Any]				= transformers.models.auto.configuration_auto.CONFIG_MAPPING

# Regex pattern used to find the checkpoint mentioned in the docstring of `config_class`.
# For example, `[bert-base-uncased](https://huggingface.co/bert-base-uncased)`
UpperCamelCase_					:						str				= re.compile(R'''\[(.+?)\]\((https://huggingface\.co/.+?)\)''')


UpperCamelCase_					:						Dict				= {
    '''DecisionTransformerConfig''',
    '''EncoderDecoderConfig''',
    '''MusicgenConfig''',
    '''RagConfig''',
    '''SpeechEncoderDecoderConfig''',
    '''TimmBackboneConfig''',
    '''VisionEncoderDecoderConfig''',
    '''VisionTextDualEncoderConfig''',
    '''LlamaConfig''',
}
def 							A_     (__a				):

					'''simple docstring'''





					A_				 =						None

					# source code of `config_class`
					A_				 =						inspect.getsource(__a				)
					A_				 =						_re_checkpoint.findall(__a				)

					# Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
					# For example, `('bert-base-uncased', 'https://huggingface.co/bert-base-uncased')`
					for ckpt_name, ckpt_link in checkpoints:
										# allow the link to end with `/`
										if ckpt_link.endswith("/"				):
															A_				 =						ckpt_link[:-1]

										# verify the checkpoint name corresponds to the checkpoint link
										A_				 =						f'https://huggingface.co/{ckpt_name}'
										if ckpt_link == ckpt_link_from_name:
															A_				 =						ckpt_name
															break

					return checkpoint
def 							A_     ():

					'''simple docstring'''





					A_				 =						[]

					for config_class in list(CONFIG_MAPPING.values()				):
										# Skip deprecated models
										if "models.deprecated" in config_class.__module__:
															continue
										A_				 =						get_checkpoint_from_config_class(__a				)

										A_				 =						config_class.__name__
										if checkpoint is None and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
															configs_without_checkpoint.append(__a				)

					if len(__a				) > 0:
										A_				 =						"\n".join(sorted(__a				)				)
										raise ValueError(f'The following configurations don\'t contain any valid checkpoint:\n{message}'				)


if __name__ == "__main__":
		check_config_docstrings_have_checkpoints()




