





'''simple docstring'''




import inspect
import unittest

from transformers import YolosConfig
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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import YolosForObjectDetection, YolosModel
 from transformers.models.yolos.modeling_yolos import YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor
class 			_lowerCAmelCase		:






     """simple docstring"""






     def __init__(   self							:			int   ,       __snake_case							:			Optional[int]   ,       __snake_case							:			List[str]=13   ,       __snake_case							:			Optional[Any]=[30, 30]   ,       __snake_case							:			Dict=2   ,       __snake_case							:			Union[str, Any]=3   ,       __snake_case							:			Any=True   ,       __snake_case							:			Optional[Any]=True   ,       __snake_case							:			int=32   ,       __snake_case							:			List[Any]=5   ,       __snake_case							:			Any=4   ,       __snake_case							:			List[Any]=37   ,       __snake_case							:			str="gelu"   ,       __snake_case							:			Optional[int]=0.1   ,       __snake_case							:			List[Any]=0.1   ,       __snake_case							:			int=10   ,       __snake_case							:			str=0.02   ,       __snake_case							:			List[Any]=3   ,       __snake_case							:			List[str]=None   ,       __snake_case							:			int=8   ,       __snake_case							:			Optional[int]=10   ,       )->	Union[str, Any]:
        snake_case 						=   parent
        snake_case 						=   batch_size
        snake_case 						=   image_size
        snake_case 						=   patch_size
        snake_case 						=   num_channels
        snake_case 						=   is_training
        snake_case 						=   use_labels
        snake_case 						=   hidden_size
        snake_case 						=   num_hidden_layers
        snake_case 						=   num_attention_heads
        snake_case 						=   intermediate_size
        snake_case 						=   hidden_act
        snake_case 						=   hidden_dropout_prob
        snake_case 						=   attention_probs_dropout_prob
        snake_case 						=   type_sequence_label_size
        snake_case 						=   initializer_range
        snake_case 						=   num_labels
        snake_case 						=   scope
        snake_case 						=   n_targets
        snake_case 						=   num_detection_tokens
        # we set the expected sequence length (which is used in several tests)
        # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token) + num_detection_tokens
        snake_case 						=   (image_size[1] // patch_size) * (image_size[0] // patch_size)
        snake_case 						=   num_patches + 1 + self.num_detection_tokens



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	Dict:
        snake_case 						=   floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]						)

        snake_case 						=   None
        if self.use_labels:
           # labels is a list of Dict (each Dict being the labels for a given example in the batch)
           snake_case 						=   []
           for i in range(self.batch_size						):
              snake_case 						=   {}
              snake_case 						=   torch.randint(
                  high=self.num_labels   ,       size=(self.n_targets,)   ,       device=__snake_case						)
              snake_case 						=   torch.rand(self.n_targets   ,       4   ,       device=__snake_case						)
              labels.append(__snake_case						)

        snake_case 						=   self.get_config()

        return config, pixel_values, labels



     def      lowerCAmelCase			(   self							:			Dict						)->	List[Any]:
        return YolosConfig(
            image_size=self.image_size   ,       patch_size=self.patch_size   ,       num_channels=self.num_channels   ,       hidden_size=self.hidden_size   ,       num_hidden_layers=self.num_hidden_layers   ,       num_attention_heads=self.num_attention_heads   ,       intermediate_size=self.intermediate_size   ,       hidden_act=self.hidden_act   ,       hidden_dropout_prob=self.hidden_dropout_prob   ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,       is_decoder=__snake_case   ,       initializer_range=self.initializer_range   ,       num_detection_tokens=self.num_detection_tokens   ,       num_labels=self.num_labels   ,       )



     def      lowerCAmelCase			(   self							:			List[Any]   ,       __snake_case							:			List[str]   ,       __snake_case							:			Tuple   ,       __snake_case							:			Optional[int]						)->	str:
        snake_case 						=   YolosModel(config=__snake_case						)
        model.to(__snake_case						)
        model.eval()
        snake_case 						=   model(__snake_case						)
        self.parent.assertEqual(
            result.last_hidden_state.shape   ,       (self.batch_size, self.expected_seq_len, self.hidden_size)						)



     def      lowerCAmelCase			(   self							:			int   ,       __snake_case							:			Optional[Any]   ,       __snake_case							:			int   ,       __snake_case							:			Dict						)->	int:
        snake_case 						=   YolosForObjectDetection(__snake_case						)
        model.to(__snake_case						)
        model.eval()

        snake_case 						=   model(pixel_values=__snake_case						)
        snake_case 						=   model(__snake_case						)

        self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.num_detection_tokens, self.num_labels + 1)						)
        self.parent.assertEqual(result.pred_boxes.shape   ,       (self.batch_size, self.num_detection_tokens, 4)						)

        snake_case 						=   model(pixel_values=__snake_case   ,       labels=__snake_case						)

        self.parent.assertEqual(result.loss.shape   ,       ()						)
        self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.num_detection_tokens, self.num_labels + 1)						)
        self.parent.assertEqual(result.pred_boxes.shape   ,       (self.batch_size, self.num_detection_tokens, 4)						)



     def      lowerCAmelCase			(   self							:			int						)->	List[str]:
        snake_case 						=   self.prepare_config_and_inputs()
        snake_case						,    snake_case						,    snake_case 						=   config_and_inputs
        snake_case 						=   {"""pixel_values""": pixel_values}
        return config, inputs_dict




@require_torch
class 			_lowerCAmelCase		(			A__							,     A__							,     unittest.TestCase	):






     """simple docstring"""






     snake_case_						      =   (YolosModel, YolosForObjectDetection) if is_torch_available() else ()
     snake_case_						      =   (
         {"feature-extraction": YolosModel, "object-detection": YolosForObjectDetection} if is_torch_available() else {}
     )

     snake_case_						      =   False
     snake_case_						      =   False
     snake_case_						      =   False
     snake_case_						      =   False



     def      lowerCAmelCase			(   self							:			str   ,       __snake_case							:			Union[str, Any]   ,       __snake_case							:			Optional[Any]   ,       __snake_case							:			Any=False						)->	Optional[Any]:
        snake_case 						=   super()._prepare_for_class(__snake_case   ,       __snake_case   ,       return_labels=__snake_case						)

        if return_labels:
           if model_class.__name__ == "YolosForObjectDetection":
              snake_case 						=   []
              for i in range(self.model_tester.batch_size						):
                 snake_case 						=   {}
                 snake_case 						=   torch.ones(
                     size=(self.model_tester.n_targets,)   ,       device=__snake_case   ,       dtype=torch.long						)
                 snake_case 						=   torch.ones(
                     self.model_tester.n_targets   ,       4   ,       device=__snake_case   ,       dtype=torch.float						)
                 labels.append(__snake_case						)
              snake_case 						=   labels

        return inputs_dict



     def      lowerCAmelCase			(   self							:			Any						)->	str:
        snake_case 						=   YolosModelTester(self						)
        snake_case 						=   ConfigTester(self   ,       config_class=__snake_case   ,       has_text_modality=__snake_case   ,       hidden_size=37						)



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	int:
        self.config_tester.run_common_tests()



     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Optional[Any]:
        # YOLOS does not use inputs_embeds
        pass



     def      lowerCAmelCase			(   self							:			Optional[int]						)->	List[str]:
        snake_case						,    snake_case 						=   self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           snake_case 						=   model_class(__snake_case						)
           self.assertIsInstance(model.get_input_embeddings()   ,       (nn.Module)						)
           snake_case 						=   model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(__snake_case   ,       nn.Linear						)						)



     def      lowerCAmelCase			(   self							:			List[Any]						)->	str:
        snake_case						,    snake_case 						=   self.model_tester.prepare_config_and_inputs_for_common()

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

           snake_case 						=   ["""pixel_values"""]
           self.assertListEqual(arg_names[:1]   ,       __snake_case						)



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	int:
        snake_case 						=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__snake_case						)



     def      lowerCAmelCase			(   self							:			Dict						)->	Optional[int]:
        snake_case						,    snake_case 						=   self.model_tester.prepare_config_and_inputs_for_common()
        snake_case 						=   True

        # in YOLOS, the seq_len is different
        snake_case 						=   self.model_tester.expected_seq_len
        for model_class in self.all_model_classes:
           snake_case 						=   True
           snake_case 						=   False
           snake_case 						=   True
           snake_case 						=   model_class(__snake_case						)
           model.to(__snake_case						)
           model.eval()
           with torch.no_grad():
              snake_case 						=   model(**self._prepare_for_class(__snake_case   ,       __snake_case						)						)
           snake_case 						=   outputs.attentions
           self.assertEqual(len(__snake_case						)   ,       self.model_tester.num_hidden_layers						)

           # check that output_attentions also work using config
           del inputs_dict["output_attentions"]
           snake_case 						=   True
           snake_case 						=   model_class(__snake_case						)
           model.to(__snake_case						)
           model.eval()
           with torch.no_grad():
              snake_case 						=   model(**self._prepare_for_class(__snake_case   ,       __snake_case						)						)
           snake_case 						=   outputs.attentions
           self.assertEqual(len(__snake_case						)   ,       self.model_tester.num_hidden_layers						)

           self.assertListEqual(
               list(attentions[0].shape[-3:]						)   ,       [self.model_tester.num_attention_heads, seq_len, seq_len]   ,       )
           snake_case 						=   len(__snake_case						)

           # Check attention is always last and order is fine
           snake_case 						=   True
           snake_case 						=   True
           snake_case 						=   model_class(__snake_case						)
           model.to(__snake_case						)
           model.eval()
           with torch.no_grad():
              snake_case 						=   model(**self._prepare_for_class(__snake_case   ,       __snake_case						)						)

           snake_case 						=   1
           self.assertEqual(out_len + added_hidden_states   ,       len(__snake_case						)						)

           snake_case 						=   outputs.attentions

           self.assertEqual(len(__snake_case						)   ,       self.model_tester.num_hidden_layers						)
           self.assertListEqual(
               list(self_attentions[0].shape[-3:]						)   ,       [self.model_tester.num_attention_heads, seq_len, seq_len]   ,       )



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	List[str]:
        def check_hidden_states_output(__snake_case							:			Optional[int]   ,       __snake_case							:			List[str]   ,       __snake_case							:			Dict						):
           snake_case 						=   model_class(__snake_case						)
           model.to(__snake_case						)
           model.eval()

           with torch.no_grad():
              snake_case 						=   model(**self._prepare_for_class(__snake_case   ,       __snake_case						)						)

           snake_case 						=   outputs.hidden_states

           snake_case 						=   getattr(
               self.model_tester   ,       """expected_num_hidden_layers"""   ,       self.model_tester.num_hidden_layers + 1						)
           self.assertEqual(len(__snake_case						)   ,       __snake_case						)

           # YOLOS has a different seq_length
           snake_case 						=   self.model_tester.expected_seq_len

           self.assertListEqual(
               list(hidden_states[0].shape[-2:]						)   ,       [seq_length, self.model_tester.hidden_size]   ,       )

        snake_case						,    snake_case 						=   self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           snake_case 						=   True
           check_hidden_states_output(__snake_case   ,       __snake_case   ,       __snake_case						)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           snake_case 						=   True

           check_hidden_states_output(__snake_case   ,       __snake_case   ,       __snake_case						)



     def      lowerCAmelCase			(   self							:			Tuple						)->	Optional[Any]:
        snake_case 						=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_object_detection(*__snake_case						)



     @slow
     def      lowerCAmelCase			(   self							:			int						)->	str:
        for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           snake_case 						=   YolosModel.from_pretrained(__snake_case						)
           self.assertIsNotNone(__snake_case						)





def 	__lowerCamelCase			(      )		->  Tuple:
   snake_case 						=   Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""		)
   return image


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






     """simple docstring"""






     @cached_property
     def      lowerCAmelCase			(   self							:			Tuple						)->	Dict:
        return AutoImageProcessor.from_pretrained("""hustvl/yolos-small"""						) if is_vision_available() else None



     @slow
     def      lowerCAmelCase			(   self							:			List[str]						)->	str:
        snake_case 						=   YolosForObjectDetection.from_pretrained("""hustvl/yolos-small"""						).to(__snake_case						)

        snake_case 						=   self.default_image_processor
        snake_case 						=   prepare_img()
        snake_case 						=   image_processor(images=__snake_case   ,       return_tensors="""pt"""						).to(__snake_case						)

        # forward pass
        with torch.no_grad():
           snake_case 						=   model(inputs.pixel_values						)

        # verify outputs
        snake_case 						=   torch.Size((1, 1_00, 92)						)
        self.assertEqual(outputs.logits.shape   ,       __snake_case						)

        snake_case 						=   torch.tensor(
            [[-24.02_48, -10.30_24, -14.82_90], [-42.03_92, -16.82_00, -27.43_34], [-27.27_43, -11.81_54, -18.71_48]]   ,       device=__snake_case   ,       )
        snake_case 						=   torch.tensor(
            [[0.25_59, 0.54_55, 0.47_06], [0.29_89, 0.72_79, 0.18_75], [0.77_32, 0.40_17, 0.44_62]]   ,       device=__snake_case						)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]   ,       __snake_case   ,       atol=1e-4						)						)
        self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3]   ,       __snake_case   ,       atol=1e-4						)						)

        # verify postprocessing
        snake_case 						=   image_processor.post_process_object_detection(
            __snake_case   ,       threshold=0.3   ,       target_sizes=[image.size[::-1]]						)[0]
        snake_case 						=   torch.tensor([0.99_94, 0.97_90, 0.99_64, 0.99_72, 0.98_61]						).to(__snake_case						)
        snake_case 						=   [75, 75, 17, 63, 17]
        snake_case 						=   torch.tensor([3_35.06_09, 79.38_48, 3_75.42_16, 1_87.24_95]						).to(__snake_case						)

        self.assertEqual(len(results["""scores"""]						)   ,       5						)
        self.assertTrue(torch.allclose(results["""scores"""]   ,       __snake_case   ,       atol=1e-4						)						)
        self.assertSequenceEqual(results["""labels"""].tolist()   ,       __snake_case						)
        self.assertTrue(torch.allclose(results["""boxes"""][0, :]   ,       __snake_case						)						)








'''simple docstring'''




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


_SCREAMING_SNAKE_CASE							     =			logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE							     =			{
    "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 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   "trajectory_transformer"
     snake_case_						      =   ["past_key_values"]
     snake_case_						      =   {
         "hidden_size": "n_embd",
         "num_attention_heads": "n_head",
         "num_hidden_layers": "n_layer",
     }



     def __init__(   self							:			Optional[Any]   ,       __snake_case							:			int=1_00   ,       __snake_case							:			List[str]=5   ,       __snake_case							:			str=1   ,       __snake_case							:			Union[str, Any]=1   ,       __snake_case							:			Union[str, Any]=2_49   ,       __snake_case							:			List[Any]=6   ,       __snake_case							:			Optional[int]=17   ,       __snake_case							:			Optional[Any]=25   ,       __snake_case							:			Union[str, Any]=4   ,       __snake_case							:			List[Any]=4   ,       __snake_case							:			Optional[int]=1_28   ,       __snake_case							:			Optional[Any]=0.1   ,       __snake_case							:			Any=0.1   ,       __snake_case							:			int=0.1   ,       __snake_case							:			Union[str, Any]=0.00_06   ,       __snake_case							:			Tuple=5_12   ,       __snake_case							:			int=0.02   ,       __snake_case							:			Any=1e-12   ,       __snake_case							:			Optional[Any]=1   ,       __snake_case							:			List[Any]=True   ,       __snake_case							:			Dict=1   ,       __snake_case							:			Dict=5_02_56   ,       __snake_case							:			Union[str, Any]=5_02_56   ,       **__snake_case							:			Tuple   ,       )->	Optional[int]:
        snake_case 						=   vocab_size
        snake_case 						=   action_weight
        snake_case 						=   reward_weight
        snake_case 						=   value_weight
        snake_case 						=   max_position_embeddings
        snake_case 						=   block_size
        snake_case 						=   action_dim
        snake_case 						=   observation_dim
        snake_case 						=   transition_dim
        snake_case 						=   learning_rate
        snake_case 						=   n_layer
        snake_case 						=   n_head
        snake_case 						=   n_embd
        snake_case 						=   embd_pdrop
        snake_case 						=   attn_pdrop
        snake_case 						=   resid_pdrop
        snake_case 						=   initializer_range
        snake_case 						=   layer_norm_eps
        snake_case 						=   kaiming_initializer_range
        snake_case 						=   use_cache
        super().__init__(pad_token_id=__snake_case   ,       bos_token_id=__snake_case   ,       eos_token_id=__snake_case   ,       **__snake_case						)

