

"""simple docstring"""

import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
					from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor






class 	_lowercase       :




	def __init__(			self		,			UpperCAmelCase_		,			UpperCAmelCase_=13		,			UpperCAmelCase_=32		,			UpperCAmelCase_=3		,			UpperCAmelCase_=4		,			UpperCAmelCase_=[10, 20, 30, 40]		,			UpperCAmelCase_=[2, 2, 3, 2]		,			UpperCAmelCase_=True		,			UpperCAmelCase_=True		,			UpperCAmelCase_=37		,			UpperCAmelCase_="gelu"		,			UpperCAmelCase_=10		,			UpperCAmelCase_=0.02		,			UpperCAmelCase_=["stage2", "stage3", "stage4"]		,			UpperCAmelCase_=[2, 3, 4]		,			UpperCAmelCase_=None		,			)		->							List[Any]:
			lowerCamelCase					:      Dict          =   parent
			lowerCamelCase					:      List[Any]          =   batch_size
			lowerCamelCase					:      List[Any]          =   image_size
			lowerCamelCase					:      Dict          =   num_channels
			lowerCamelCase					:      Union[str, Any]          =   num_stages
			lowerCamelCase					:      List[str]          =   hidden_sizes
			lowerCamelCase					:      Tuple          =   depths
			lowerCamelCase					:      Tuple          =   is_training
			lowerCamelCase					:      List[str]          =   use_labels
			lowerCamelCase					:      Any          =   intermediate_size
			lowerCamelCase					:      List[str]          =   hidden_act
			lowerCamelCase					:      Tuple          =   num_labels
			lowerCamelCase					:      Tuple          =   initializer_range
			lowerCamelCase					:      List[Any]          =   out_features
			lowerCamelCase					:      int          =   out_indices
			lowerCamelCase					:      Optional[Any]          =   scope




	def       _UpperCamelCase   (			self )		->							Dict:
			lowerCamelCase					:      str          =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

			lowerCamelCase					:      int          =   None
			if self.use_labels:
					lowerCamelCase					:      int          =   ids_tensor([self.batch_size]		,			self.num_labels )

			lowerCamelCase					:      Tuple          =   self.get_config()

			return config, pixel_values, labels




	def       _UpperCamelCase   (			self )		->							Tuple:
			return ConvNextVaConfig(
			    num_channels=self.num_channels		,			hidden_sizes=self.hidden_sizes		,			depths=self.depths		,			num_stages=self.num_stages		,			hidden_act=self.hidden_act		,			is_decoder=_A		,			initializer_range=self.initializer_range		,			out_features=self.out_features		,			out_indices=self.out_indices		,			num_labels=self.num_labels		,			)




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							Dict:
			lowerCamelCase					:      int          =   ConvNextVaModel(config=_A )
			model.to(_A )
			model.eval()
			lowerCamelCase					:      str          =   model(_A )
			# expected last hidden states: B, C, H // 32, W // 32
			self.parent.assertEqual(
			    result.last_hidden_state.shape		,			(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)		,			)




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							Any:
			lowerCamelCase					:      Union[str, Any]          =   ConvNextVaForImageClassification(_A )
			model.to(_A )
			model.eval()
			lowerCamelCase					:      Optional[Any]          =   model(_A		,			labels=_A )
			self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.num_labels) )




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							Dict:
			lowerCamelCase					:      Tuple          =   ConvNextVaBackbone(config=_A )
			model.to(_A )
			model.eval()
			lowerCamelCase					:      Optional[Any]          =   model(_A )

			# verify hidden states
			self.parent.assertEqual(len(result.feature_maps )		,			len(config.out_features ) )
			self.parent.assertListEqual(list(result.feature_maps[0].shape )		,			[self.batch_size, self.hidden_sizes[1], 4, 4] )

			# verify channels
			self.parent.assertEqual(len(model.channels )		,			len(config.out_features ) )
			self.parent.assertListEqual(model.channels		,			config.hidden_sizes[1:] )

			# verify backbone works with out_features=None
			lowerCamelCase					:      List[str]          =   None
			lowerCamelCase					:      Union[str, Any]          =   ConvNextVaBackbone(config=_A )
			model.to(_A )
			model.eval()
			lowerCamelCase					:      str          =   model(_A )

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps )		,			1 )
			self.parent.assertListEqual(list(result.feature_maps[0].shape )		,			[self.batch_size, self.hidden_sizes[-1], 1, 1] )

			# verify channels
			self.parent.assertEqual(len(model.channels )		,			1 )
			self.parent.assertListEqual(model.channels		,			[config.hidden_sizes[-1]] )




	def       _UpperCamelCase   (			self )		->							Optional[int]:
			lowerCamelCase					:      List[str]          =   self.prepare_config_and_inputs()
			lowerCamelCase     ,   lowerCamelCase     ,   lowerCamelCase					:      Optional[Any]          =   config_and_inputs
			lowerCamelCase					:      Union[str, Any]          =   {'pixel_values': pixel_values}
			return config, inputs_dict




	def       _UpperCamelCase   (			self )		->							Optional[int]:
			lowerCamelCase					:      Any          =   self.prepare_config_and_inputs()
			lowerCamelCase     ,   lowerCamelCase     ,   lowerCamelCase					:      str          =   config_and_inputs
			lowerCamelCase					:      Optional[int]          =   {'pixel_values': pixel_values, 'labels': labels}
			return config, inputs_dict






@require_torch
class 	_lowercase       (    snake_case__    ,	snake_case__    ,	unittest.TestCase		):
	lowercase_				=     (
	    (
	        ConvNextVaModel,
	        ConvNextVaForImageClassification,
	        ConvNextVaBackbone,
	    )
	    if is_torch_available()
	    else ()
	)
	lowercase_				=     (
	    {"""feature-extraction""": ConvNextVaModel, """image-classification""": ConvNextVaForImageClassification}
	    if is_torch_available()
	    else {}
	)

	lowercase_				=     False
	lowercase_				=     False
	lowercase_				=     False
	lowercase_				=     False
	lowercase_				=     False




	def       _UpperCamelCase   (			self )		->							Dict:
			lowerCamelCase					:      str          =   ConvNextVaModelTester(self )
			lowerCamelCase					:      Optional[int]          =   ConfigTester(self		,			config_class=_A		,			has_text_modality=_A		,			hidden_size=37 )




	def       _UpperCamelCase   (			self )		->							Union[str, Any]:
			self.create_and_test_config_common_properties()
			self.config_tester.create_and_test_config_to_json_string()
			self.config_tester.create_and_test_config_to_json_file()
			self.config_tester.create_and_test_config_from_and_save_pretrained()
			self.config_tester.create_and_test_config_with_num_labels()
			self.config_tester.check_config_can_be_init_without_params()
			self.config_tester.check_config_arguments_init()




	def       _UpperCamelCase   (			self )		->							List[Any]:
			return




	@unittest.skip(reason='ConvNextV2 does not use inputs_embeds' )
	def       _UpperCamelCase   (			self )		->							str:
			pass




	@unittest.skip(reason='ConvNextV2 does not support input and output embeddings' )
	def       _UpperCamelCase   (			self )		->							str:
			pass




	@unittest.skip(reason='ConvNextV2 does not use feedforward chunking' )
	def       _UpperCamelCase   (			self )		->							Union[str, Any]:
			pass




	def       _UpperCamelCase   (			self )		->							str:



			if not self.model_tester.is_training:
					return

			for model_class in self.all_model_classes:
					lowerCamelCase     ,   lowerCamelCase					:      Tuple          =   self.model_tester.prepare_config_and_inputs_with_labels()
					lowerCamelCase					:      Optional[Any]          =   True

					if model_class.__name__ in [
					    *get_values(_A ),
					    *get_values(_A ),
					]:
							continue

					lowerCamelCase					:      int          =   model_class(_A )
					model.to(_A )
					model.train()
					lowerCamelCase					:      List[Any]          =   self._prepare_for_class(_A		,			_A		,			return_labels=_A )
					lowerCamelCase					:      List[Any]          =   model(**_A ).loss
					loss.backward()




	def       _UpperCamelCase   (			self )		->							Optional[int]:



			if not self.model_tester.is_training:
					return

			for model_class in self.all_model_classes:
					lowerCamelCase     ,   lowerCamelCase					:      Optional[int]          =   self.model_tester.prepare_config_and_inputs_with_labels()
					lowerCamelCase					:      List[str]          =   False
					lowerCamelCase					:      Dict          =   True

					if (
					    model_class.__name__
					    in [*get_values(_A ), *get_values(_A )]
					    or not model_class.supports_gradient_checkpointing
					):
							continue

					lowerCamelCase					:      int          =   model_class(_A )
					model.to(_A )
					model.gradient_checkpointing_enable()
					model.train()
					lowerCamelCase					:      int          =   self._prepare_for_class(_A		,			_A		,			return_labels=_A )
					lowerCamelCase					:      Optional[Any]          =   model(**_A ).loss
					loss.backward()




	def       _UpperCamelCase   (			self )		->							int:
			lowerCamelCase     ,   lowerCamelCase					:      Dict          =   self.model_tester.prepare_config_and_inputs_for_common()

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

					lowerCamelCase					:      Optional[int]          =   ['pixel_values']
					self.assertListEqual(arg_names[:1]		,			_A )




	def       _UpperCamelCase   (			self )		->							Union[str, Any]:
			lowerCamelCase					:      str          =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*_A )




	def       _UpperCamelCase   (			self )		->							List[str]:



			def check_hidden_states_output(UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ ):
					lowerCamelCase					:      Dict          =   model_class(_A )
					model.to(_A )
					model.eval()

					with torch.no_grad():
							lowerCamelCase					:      str          =   model(**self._prepare_for_class(_A		,			_A ) )

					lowerCamelCase					:      Optional[int]          =   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

					lowerCamelCase					:      Any          =   self.model_tester.num_stages
					self.assertEqual(len(_A )		,			expected_num_stages + 1 )

					# ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
					self.assertListEqual(
					    list(hidden_states[0].shape[-2:] )		,			[self.model_tester.image_size // 4, self.model_tester.image_size // 4]		,			)

			lowerCamelCase     ,   lowerCamelCase					:      int          =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					lowerCamelCase					:      Any          =   True
					check_hidden_states_output(_A		,			_A		,			_A )

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					lowerCamelCase					:      Optional[int]          =   True

					check_hidden_states_output(_A		,			_A		,			_A )




	def       _UpperCamelCase   (			self )		->							Tuple:
			lowerCamelCase					:      str          =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*_A )




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



			for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase					:      List[str]          =   ConvNextVaModel.from_pretrained(_A )
					self.assertIsNotNone(_A )


def 		UpperCAmelCase (		):





		'''simple docstring'''


		lowerCamelCase					:      Optional[int]          =   Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
		return image






@require_torch
@require_vision
class 	_lowercase       (    unittest.TestCase		):




	@cached_property
	def       _UpperCamelCase   (			self )		->							Any:
			return AutoImageProcessor.from_pretrained('facebook/convnextv2-tiny-1k-224' ) if is_vision_available() else None




	@slow
	def       _UpperCamelCase   (			self )		->							Optional[int]:
			lowerCamelCase					:      List[Any]          =   ConvNextVaForImageClassification.from_pretrained('facebook/convnextv2-tiny-1k-224' ).to(_A )

			lowerCamelCase					:      int          =   self.default_image_processor
			lowerCamelCase					:      int          =   prepare_img()
			lowerCamelCase					:      Tuple          =   preprocessor(images=_A		,			return_tensors='pt' ).to(_A )

			# forward pass
			with torch.no_grad():
					lowerCamelCase					:      List[Any]          =   model(**_A )

			# verify the logits
			lowerCamelCase					:      List[Any]          =   torch.Size((1, 1000) )
			self.assertEqual(outputs.logits.shape		,			_A )

			lowerCamelCase					:      Optional[Any]          =   torch.tensor([0.9996, 0.1966, -0.4386] ).to(_A )
			self.assertTrue(torch.allclose(outputs.logits[0, :3]		,			_A		,			atol=1E-4 ) )








"""simple docstring"""

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


_A						=			logging.get_logger(__name__)

_A						=			{
    'CarlCochet/trajectory-transformer-halfcheetah-medium-v2': (
        'https://huggingface.co/CarlCochet/trajectory-transformer-halfcheetah-medium-v2/resolve/main/config.json'
    ),
    # See all TrajectoryTransformer models at https://huggingface.co/models?filter=trajectory_transformer
}






class 	_lowercase       (    __UpperCAmelCase		):
	lowercase_				=     'trajectory_transformer'
	lowercase_				=     ['past_key_values']
	lowercase_				=     {
	    'hidden_size': 'n_embd',
	    'num_attention_heads': 'n_head',
	    'num_hidden_layers': 'n_layer',
	}




	def __init__(			self		,			UpperCAmelCase_=100		,			UpperCAmelCase_=5		,			UpperCAmelCase_=1		,			UpperCAmelCase_=1		,			UpperCAmelCase_=249		,			UpperCAmelCase_=6		,			UpperCAmelCase_=17		,			UpperCAmelCase_=25		,			UpperCAmelCase_=4		,			UpperCAmelCase_=4		,			UpperCAmelCase_=128		,			UpperCAmelCase_=0.1		,			UpperCAmelCase_=0.1		,			UpperCAmelCase_=0.1		,			UpperCAmelCase_=0.0006		,			UpperCAmelCase_=512		,			UpperCAmelCase_=0.02		,			UpperCAmelCase_=1E-1_2		,			UpperCAmelCase_=1		,			UpperCAmelCase_=True		,			UpperCAmelCase_=1		,			UpperCAmelCase_=50256		,			UpperCAmelCase_=50256		,			**UpperCAmelCase_		,			)		->							List[Any]:
			lowerCamelCase					:      int          =   vocab_size
			lowerCamelCase					:      List[str]          =   action_weight
			lowerCamelCase					:      List[Any]          =   reward_weight
			lowerCamelCase					:      List[str]          =   value_weight
			lowerCamelCase					:      Tuple          =   max_position_embeddings
			lowerCamelCase					:      List[str]          =   block_size
			lowerCamelCase					:      Any          =   action_dim
			lowerCamelCase					:      List[Any]          =   observation_dim
			lowerCamelCase					:      Any          =   transition_dim
			lowerCamelCase					:      int          =   learning_rate
			lowerCamelCase					:      Union[str, Any]          =   n_layer
			lowerCamelCase					:      Tuple          =   n_head
			lowerCamelCase					:      Any          =   n_embd
			lowerCamelCase					:      Union[str, Any]          =   embd_pdrop
			lowerCamelCase					:      Optional[int]          =   attn_pdrop
			lowerCamelCase					:      int          =   resid_pdrop
			lowerCamelCase					:      Optional[int]          =   initializer_range
			lowerCamelCase					:      List[str]          =   layer_norm_eps
			lowerCamelCase					:      Any          =   kaiming_initializer_range
			lowerCamelCase					:      str          =   use_cache
			super().__init__(pad_token_id=UpperCAmelCase_		,			bos_token_id=UpperCAmelCase_		,			eos_token_id=UpperCAmelCase_		,			**UpperCAmelCase_ )






