





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






				"""simple docstring"""







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




				def   UpperCamelCase__   (  self     :					List[str]					):
								_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=10_00					)
								_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   UpperCamelCase__   (  self     :					List[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     :					Dict					,				__a     :					List[str]					,				__a     :					Any					,				__a     :					Tuple					,				__a     :					Tuple					,				__a     :					Optional[int]					,				__a     :					str					,				__a     :					List[Any]					,				):
								_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   UpperCamelCase__   (  self     :					Optional[int]					):
								_a								= self.prepare_config_and_inputs()
								(
								    (
								    _a
								)	,					(
								    _a
								)	,					(
								    _a
								)	,					(
								    _a
								)	,					(
								    _a
								)	,					(
								    _a
								)	,					(
								    _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 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_       ,							lowerCamelCase_       ,							lowerCamelCase_       ,							unittest.TestCase      ):






				"""simple docstring"""







				__a            =(DecisionTransformerModel,) if is_torch_available() else ()
				__a            =()
				__a            ={'feature-extraction': DecisionTransformerModel} if is_torch_available() else {}

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

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




				def   UpperCamelCase__   (  self     :					int					):
								_a								= DecisionTransformerModelTester(self					)
								_a								= ConfigTester(self					,				config_class=__a					,				hidden_size=37					)




				def   UpperCamelCase__   (  self     :					Any					):
								self.config_tester.run_common_tests()




				def   UpperCamelCase__   (  self     :					List[Any]					):
								_a								= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__a					)




				@slow
				def   UpperCamelCase__   (  self     :					Optional[int]					):
								for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_a								= DecisionTransformerModel.from_pretrained(__a					)
												self.assertIsNotNone(__a					)




				def   UpperCamelCase__   (  self     :					List[Any]					):
								_a	,					_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 				__SCREAMING_SNAKE_CASE  (unittest.TestCase      ):






				"""simple docstring"""







				@slow
				def   UpperCamelCase__   (  self     :					List[Any]					):
								_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.242793, -0.28693074, 0.8742613], [0.67815274, -0.08101085, -0.12952147]]					,				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	,					_a	,					_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	,					_a	,					_a	,					_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 unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible
class 				__SCREAMING_SNAKE_CASE  (unittest.TestCase      ):






				"""simple docstring"""







				def   UpperCamelCase__   (  self     :					str					):
								_a								= [
								    "safety_checker/pytorch_model.bin",
								    "safety_checker/model.safetensors",
								    "vae/diffusion_pytorch_model.bin",
								    "vae/diffusion_pytorch_model.safetensors",
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								self.assertTrue(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					List[str]					):
								_a								= [
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								self.assertTrue(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					List[str]					):
								_a								= [
								    "safety_checker/pytorch_model.bin",
								    "safety_checker/model.safetensors",
								    "vae/diffusion_pytorch_model.bin",
								    "vae/diffusion_pytorch_model.safetensors",
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								    "unet/diffusion_pytorch_model.bin",
								    # Removed: 'unet/diffusion_pytorch_model.safetensors',
								]
								self.assertFalse(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					List[str]					):
								_a								= [
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								]
								self.assertTrue(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					Optional[Any]					):
								_a								= [
								    "safety_checker/pytorch_model.bin",
								    "safety_checker/model.safetensors",
								    "vae/diffusion_pytorch_model.bin",
								    "vae/diffusion_pytorch_model.safetensors",
								    "text_encoder/pytorch_model.bin",
								    # Removed: 'text_encoder/model.safetensors',
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								self.assertFalse(is_safetensors_compatible(__a					)					)




				def   UpperCamelCase__   (  self     :					str					):
								_a								= [
								    "safety_checker/pytorch_model.fp16.bin",
								    "safety_checker/model.fp16.safetensors",
								    "vae/diffusion_pytorch_model.fp16.bin",
								    "vae/diffusion_pytorch_model.fp16.safetensors",
								    "text_encoder/pytorch_model.fp16.bin",
								    "text_encoder/model.fp16.safetensors",
								    "unet/diffusion_pytorch_model.fp16.bin",
								    "unet/diffusion_pytorch_model.fp16.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Any					):
								_a								= [
								    "unet/diffusion_pytorch_model.fp16.bin",
								    "unet/diffusion_pytorch_model.fp16.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Any					):
								# pass variant but use the non-variant filenames
								_a								= [
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Optional[Any]					):
								_a								= [
								    "safety_checker/pytorch_model.fp16.bin",
								    "safety_checker/model.fp16.safetensors",
								    "vae/diffusion_pytorch_model.fp16.bin",
								    "vae/diffusion_pytorch_model.fp16.safetensors",
								    "text_encoder/pytorch_model.fp16.bin",
								    "text_encoder/model.fp16.safetensors",
								    "unet/diffusion_pytorch_model.fp16.bin",
								    # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
								]
								_a								= "fp16"
								self.assertFalse(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Dict					):
								_a								= [
								    "text_encoder/pytorch_model.fp16.bin",
								    "text_encoder/model.fp16.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					List[str]					):
								# pass variant but use the non-variant filenames
								_a								= [
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								]
								_a								= "fp16"
								self.assertTrue(is_safetensors_compatible(__a					,				variant=__a					)					)




				def   UpperCamelCase__   (  self     :					Optional[int]					):
								_a								= [
								    "safety_checker/pytorch_model.fp16.bin",
								    "safety_checker/model.fp16.safetensors",
								    "vae/diffusion_pytorch_model.fp16.bin",
								    "vae/diffusion_pytorch_model.fp16.safetensors",
								    "text_encoder/pytorch_model.fp16.bin",
								    # 'text_encoder/model.fp16.safetensors',
								    "unet/diffusion_pytorch_model.fp16.bin",
								    "unet/diffusion_pytorch_model.fp16.safetensors",
								]
								_a								= "fp16"
								self.assertFalse(is_safetensors_compatible(__a					,				variant=__a					)					)
