





import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, 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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch
       from torch import nn

       from transformers import (
           MODEL_MAPPING,
           BeitForImageClassification,
           BeitForMaskedImageModeling,
           BeitForSemanticSegmentation,
           BeitModel,
       )
       from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       import PIL
       from PIL import Image

       from transformers import BeitImageProcessor





class 			lowerCAmelCase__:






       '''simple docstring'''


       def __init__(      self    ,						__lowerCamelCase    ,						__lowerCamelCase=1_0_0    ,						__lowerCamelCase=1_3    ,						__lowerCamelCase=3_0    ,						__lowerCamelCase=2    ,						__lowerCamelCase=3    ,						__lowerCamelCase=True    ,						__lowerCamelCase=True    ,						__lowerCamelCase=3_2    ,						__lowerCamelCase=4    ,						__lowerCamelCase=4    ,						__lowerCamelCase=3_7    ,						__lowerCamelCase="gelu"    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=1_0    ,						__lowerCamelCase=0.02    ,						__lowerCamelCase=3    ,						__lowerCamelCase=None    ,						__lowerCamelCase=[0, 1, 2, 3]    ,						)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							parent
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							1_0_0
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							batch_size
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							image_size
            _SCREAMING_SNAKE_CASE       :	Dict						    =							patch_size
            _SCREAMING_SNAKE_CASE       :	Any						    =							num_channels
            _SCREAMING_SNAKE_CASE       :	str						    =							is_training
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							use_labels
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							hidden_size
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							num_hidden_layers
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							num_attention_heads
            _SCREAMING_SNAKE_CASE       :	Any						    =							intermediate_size
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							hidden_act
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							hidden_dropout_prob
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							attention_probs_dropout_prob
            _SCREAMING_SNAKE_CASE       :	Any						    =							type_sequence_label_size
            _SCREAMING_SNAKE_CASE       :	Any						    =							initializer_range
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							scope
            _SCREAMING_SNAKE_CASE       :	int						    =							out_indices
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							num_labels

            # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							(image_size // patch_size) ** 2
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							num_patches + 1




       def     UpperCamelCase_			(      self							)			->     Optional[int]:
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							None
            _SCREAMING_SNAKE_CASE       :	str						    =							None
            if self.use_labels:
                 _SCREAMING_SNAKE_CASE       :	List[str]						    =							ids_tensor([self.batch_size]    ,						self.type_sequence_label_size							)
                 _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							ids_tensor([self.batch_size, self.image_size, self.image_size]    ,						self.num_labels							)

            _SCREAMING_SNAKE_CASE       :	Tuple						    =							self.get_config()

            return config, pixel_values, labels, pixel_labels




       def     UpperCamelCase_			(      self							)			->     Union[str, Any]:
            return BeitConfig(
                vocab_size=self.vocab_size    ,						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=__lowerCamelCase    ,						initializer_range=self.initializer_range    ,						out_indices=self.out_indices    ,						)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     List[str]:
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							BeitModel(config=__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	int						    =							model(__lowerCamelCase							)
            self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)							)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     str:
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							BeitForMaskedImageModeling(config=__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							model(__lowerCamelCase							)
            self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length - 1, self.vocab_size)							)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     List[str]:
            _SCREAMING_SNAKE_CASE       :	str						    =							self.type_sequence_label_size
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							BeitForImageClassification(__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							model(__lowerCamelCase    ,						labels=__lowerCamelCase							)
            self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.type_sequence_label_size)							)

            # test greyscale images
            _SCREAMING_SNAKE_CASE       :	Any						    =							1
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							BeitForImageClassification(__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()

            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							model(__lowerCamelCase    ,						labels=__lowerCamelCase							)
            self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.type_sequence_label_size)							)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     List[Any]:
            _SCREAMING_SNAKE_CASE       :	Dict						    =							self.num_labels
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							BeitForSemanticSegmentation(__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							model(__lowerCamelCase							)
            self.parent.assertEqual(
                result.logits.shape    ,						(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)							)
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							model(__lowerCamelCase    ,						labels=__lowerCamelCase							)
            self.parent.assertEqual(
                result.logits.shape    ,						(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)							)







       def     UpperCamelCase_			(      self							)			->     Dict:
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.prepare_config_and_inputs()
            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	List[str]						    =							config_and_inputs
            _SCREAMING_SNAKE_CASE       :	Any						    =							{"pixel_values": pixel_values}
            return config, inputs_dict





@require_torch
class 			lowerCAmelCase__( __lowercase  ,     __lowercase  ,     unittest.TestCase    ):






       '''simple docstring'''


       __snake_case 			=    (
           (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
           if is_torch_available()
           else ()
       )
       __snake_case 			=    (
           {
               'feature-extraction': BeitModel,
               'image-classification': BeitForImageClassification,
               'image-segmentation': BeitForSemanticSegmentation,
           }
           if is_torch_available()
           else {}
       )

       __snake_case 			=    False
       __snake_case 			=    False
       __snake_case 			=    False




       def     UpperCamelCase_			(      self							)			->     List[Any]:
            _SCREAMING_SNAKE_CASE       :	int						    =							BeitModelTester(self							)
            _SCREAMING_SNAKE_CASE       :	int						    =							ConfigTester(self    ,						config_class=__lowerCamelCase    ,						has_text_modality=__lowerCamelCase    ,						hidden_size=3_7							)




       def     UpperCamelCase_			(      self							)			->     Union[str, Any]:
            self.config_tester.run_common_tests()




       @unittest.skip(reason="BEiT does not use inputs_embeds"							)
       def     UpperCamelCase_			(      self							)			->     Tuple:
            pass




       @require_torch_multi_gpu
       @unittest.skip(reason="BEiT has some layers using `add_module` which doesn't work well with `nn.DataParallel`"							)
       def     UpperCamelCase_			(      self							)			->     List[Any]:
            pass




       def     UpperCamelCase_			(      self							)			->     Dict:
            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Dict						    =							self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 _SCREAMING_SNAKE_CASE       :	int						    =							model_class(__lowerCamelCase							)
                 self.assertIsInstance(model.get_input_embeddings()    ,						(nn.Module)							)
                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							model.get_output_embeddings()
                 self.assertTrue(x is None or isinstance(__lowerCamelCase    ,						nn.Linear							)							)




       def     UpperCamelCase_			(      self							)			->     Any:
            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.model_tester.prepare_config_and_inputs_for_common()

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

                 _SCREAMING_SNAKE_CASE       :	Tuple						    =							["pixel_values"]
                 self.assertListEqual(arg_names[:1]    ,						__lowerCamelCase							)




       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*__lowerCamelCase							)




       def     UpperCamelCase_			(      self							)			->     List[str]:
            _SCREAMING_SNAKE_CASE       :	str						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase							)




       def     UpperCamelCase_			(      self							)			->     Dict:
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase							)




       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_semantic_segmentation(*__lowerCamelCase							)




       def     UpperCamelCase_			(      self							)			->     str:
            if not self.model_tester.is_training:
                 return

            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.model_tester.prepare_config_and_inputs_for_common()
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							True

            for model_class in self.all_model_classes:
                 # we don't test BeitForMaskedImageModeling
                 if model_class in [*get_values(__lowerCamelCase							), BeitForMaskedImageModeling]:
                      continue

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




       def     UpperCamelCase_			(      self							)			->     Dict:
            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	str						    =							self.model_tester.prepare_config_and_inputs_for_common()
            if not self.model_tester.is_training:
                 return

            _SCREAMING_SNAKE_CASE       :	int						    =							False
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							True

            for model_class in self.all_model_classes:
                 # we don't test BeitForMaskedImageModeling
                 if (
                     model_class in [*get_values(__lowerCamelCase							), BeitForMaskedImageModeling]
                     or not model_class.supports_gradient_checkpointing
                 ):
                      continue

                 _SCREAMING_SNAKE_CASE       :	List[Any]						    =							model_class(__lowerCamelCase							)
                 model.gradient_checkpointing_enable()
                 model.to(__lowerCamelCase							)
                 model.train()
                 _SCREAMING_SNAKE_CASE       :	List[Any]						    =							self._prepare_for_class(__lowerCamelCase    ,						__lowerCamelCase    ,						return_labels=__lowerCamelCase							)
                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							model(**__lowerCamelCase							).loss
                 loss.backward()




       def     UpperCamelCase_			(      self							)			->     Tuple:
            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	List[str]						    =							self.model_tester.prepare_config_and_inputs_for_common()

            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							_config_zero_init(__lowerCamelCase							)
            for model_class in self.all_model_classes:
                 _SCREAMING_SNAKE_CASE       :	str						    =							model_class(config=__lowerCamelCase							)
                 for name, param in model.named_parameters():
                      # we skip lambda parameters as these require special initial values
                      # determined by config.layer_scale_init_value
                      if "lambda" in name:
                           continue
                      if param.requires_grad:
                           self.assertIn(
                               ((param.data.mean() * 1E9).round() / 1E9).item()    ,						[0.0, 1.0]    ,						msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""    ,						)







       @slow
       def     UpperCamelCase_			(      self							)			->     Optional[int]:
            for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _SCREAMING_SNAKE_CASE       :	str						    =							BeitModel.from_pretrained(__lowerCamelCase							)
                 self.assertIsNotNone(__lowerCamelCase							)






def     lowerCamelCase__			():
     _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
     return image





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






       '''simple docstring'''


       @cached_property
       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"							) if is_vision_available() else None




       @slow
       def     UpperCamelCase_			(      self							)			->     Tuple:
            _SCREAMING_SNAKE_CASE       :	Dict						    =							BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"							).to(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.default_image_processor
            _SCREAMING_SNAKE_CASE       :	int						    =							prepare_img()
            _SCREAMING_SNAKE_CASE       :	Any						    =							image_processor(images=__lowerCamelCase    ,						return_tensors="pt"							).pixel_values.to(__lowerCamelCase							)

            # prepare bool_masked_pos
            _SCREAMING_SNAKE_CASE       :	str						    =							torch.ones((1, 1_9_6)    ,						dtype=torch.bool							).to(__lowerCamelCase							)

            # forward pass
            with torch.no_grad():
                 _SCREAMING_SNAKE_CASE       :	List[Any]						    =							model(pixel_values=__lowerCamelCase    ,						bool_masked_pos=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	int						    =							outputs.logits

            # verify the logits
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							torch.Size((1, 1_9_6, 8_1_9_2)							)
            self.assertEqual(logits.shape    ,						__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Tuple						    =							torch.tensor(
                [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]							).to(__lowerCamelCase							)

            self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]    ,						__lowerCamelCase    ,						atol=1E-2							)							)




       @slow
       def     UpperCamelCase_			(      self							)			->     List[str]:
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"							).to(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	int						    =							self.default_image_processor
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							prepare_img()
            _SCREAMING_SNAKE_CASE       :	Dict						    =							image_processor(images=__lowerCamelCase    ,						return_tensors="pt"							).to(__lowerCamelCase							)

            # forward pass
            with torch.no_grad():
                 _SCREAMING_SNAKE_CASE       :	Dict						    =							model(**__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Dict						    =							outputs.logits

            # verify the logits
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							torch.Size((1, 1_0_0_0)							)
            self.assertEqual(logits.shape    ,						__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							torch.tensor([-1.2385, -1.0987, -1.0108]							).to(__lowerCamelCase							)

            self.assertTrue(torch.allclose(logits[0, :3]    ,						__lowerCamelCase    ,						atol=1E-4							)							)

            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							2_8_1
            self.assertEqual(logits.argmax(-1							).item()    ,						__lowerCamelCase							)




       @slow
       def     UpperCamelCase_			(      self							)			->     Union[str, Any]:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							BeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"							).to(
                __lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							self.default_image_processor
            _SCREAMING_SNAKE_CASE       :	str						    =							prepare_img()
            _SCREAMING_SNAKE_CASE       :	str						    =							image_processor(images=__lowerCamelCase    ,						return_tensors="pt"							).to(__lowerCamelCase							)

            # forward pass
            with torch.no_grad():
                 _SCREAMING_SNAKE_CASE       :	Dict						    =							model(**__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							outputs.logits

            # verify the logits
            _SCREAMING_SNAKE_CASE       :	str						    =							torch.Size((1, 2_1_8_4_1)							)
            self.assertEqual(logits.shape    ,						__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							torch.tensor([1.6881, -0.2787, 0.5901]							).to(__lowerCamelCase							)

            self.assertTrue(torch.allclose(logits[0, :3]    ,						__lowerCamelCase    ,						atol=1E-4							)							)

            _SCREAMING_SNAKE_CASE       :	Tuple						    =							2_3_9_6
            self.assertEqual(logits.argmax(-1							).item()    ,						__lowerCamelCase							)




       @slow
       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"							)
            _SCREAMING_SNAKE_CASE       :	str						    =							model.to(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							BeitImageProcessor(do_resize=__lowerCamelCase    ,						size=6_4_0    ,						do_center_crop=__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							load_dataset("hf-internal-testing/fixtures_ade20k"    ,						split="test"							)
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							Image.open(ds[0]["file"]							)
            _SCREAMING_SNAKE_CASE       :	Any						    =							image_processor(images=__lowerCamelCase    ,						return_tensors="pt"							).to(__lowerCamelCase							)

            # forward pass
            with torch.no_grad():
                 _SCREAMING_SNAKE_CASE       :	Any						    =							model(**__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							outputs.logits

            # verify the logits
            _SCREAMING_SNAKE_CASE       :	str						    =							torch.Size((1, 1_5_0, 1_6_0, 1_6_0)							)
            self.assertEqual(logits.shape    ,						__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							version.parse(PIL.__version__							) < version.parse("9.0.0"							)

            if is_pillow_less_than_a:
                 _SCREAMING_SNAKE_CASE       :	Tuple						    =							torch.tensor(
                     [
                         [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
                         [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
                         [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
                     ]    ,						device=__lowerCamelCase    ,						)
            else:
                 _SCREAMING_SNAKE_CASE       :	List[str]						    =							torch.tensor(
                     [
                         [[-4.8960, -2.3688, -3.0355], [-2.8478, -0.9836, -1.7418], [-2.9449, -1.3332, -2.1456]],
                         [[-5.8081, -3.4124, -4.1006], [-3.8561, -2.2081, -3.0323], [-3.8365, -2.4601, -3.3669]],
                         [[-0.0309, 3.9868, 4.0540], [2.9640, 4.6877, 4.9976], [3.2081, 4.7690, 4.9942]],
                     ]    ,						device=__lowerCamelCase    ,						)

            self.assertTrue(torch.allclose(logits[0, :3, :3, :3]    ,						__lowerCamelCase    ,						atol=1E-4							)							)







       @slow
       def     UpperCamelCase_			(      self							)			->     int:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"							)
            _SCREAMING_SNAKE_CASE       :	Any						    =							model.to(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							BeitImageProcessor(do_resize=__lowerCamelCase    ,						size=6_4_0    ,						do_center_crop=__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	int						    =							load_dataset("hf-internal-testing/fixtures_ade20k"    ,						split="test"							)
            _SCREAMING_SNAKE_CASE       :	str						    =							Image.open(ds[0]["file"]							)
            _SCREAMING_SNAKE_CASE       :	Dict						    =							image_processor(images=__lowerCamelCase    ,						return_tensors="pt"							).to(__lowerCamelCase							)

            # forward pass
            with torch.no_grad():
                 _SCREAMING_SNAKE_CASE       :	int						    =							model(**__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							outputs.logits.detach().cpu()

            _SCREAMING_SNAKE_CASE       :	Any						    =							image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase    ,						target_sizes=[(5_0_0, 3_0_0)]							)
            _SCREAMING_SNAKE_CASE       :	Dict						    =							torch.Size((5_0_0, 3_0_0)							)
            self.assertEqual(segmentation[0].shape    ,						__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Any						    =							torch.Size((1_6_0, 1_6_0)							)
            self.assertEqual(segmentation[0].shape    ,						__lowerCamelCase							)






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    ,						__lowerCamelCase    ,						__lowerCamelCase=1_3    ,						__lowerCamelCase=7    ,						__lowerCamelCase=6    ,						__lowerCamelCase=1_7    ,						__lowerCamelCase=2_3    ,						__lowerCamelCase=1_1    ,						__lowerCamelCase=True    ,						)			->     Optional[int]:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							parent
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							batch_size
            _SCREAMING_SNAKE_CASE       :	int						    =							seq_length
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							act_dim
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							state_dim
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							hidden_size
            _SCREAMING_SNAKE_CASE       :	Any						    =							max_length
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							is_training




       def     UpperCamelCase_			(      self							)			->     List[Any]:
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							floats_tensor((self.batch_size, self.seq_length, self.state_dim)							)
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							floats_tensor((self.batch_size, self.seq_length, self.act_dim)							)
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							floats_tensor((self.batch_size, self.seq_length, 1)							)
            _SCREAMING_SNAKE_CASE       :	str						    =							floats_tensor((self.batch_size, self.seq_length, 1)							)
            _SCREAMING_SNAKE_CASE       :	str						    =							ids_tensor((self.batch_size, self.seq_length)    ,						vocab_size=1_0_0_0							)
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							random_attention_mask((self.batch_size, self.seq_length)							)

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							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    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						)			->     Any:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							DecisionTransformerModel(config=__lowerCamelCase							)
            model.to(__lowerCamelCase							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							model(__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)

            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							)			->     Dict:
            _SCREAMING_SNAKE_CASE       :	int						    =							self.prepare_config_and_inputs()
            (
                (
                _SCREAMING_SNAKE_CASE
            )    ,			(
                _SCREAMING_SNAKE_CASE
            )    ,			(
                _SCREAMING_SNAKE_CASE
            )    ,			(
                _SCREAMING_SNAKE_CASE
            )    ,			(
                _SCREAMING_SNAKE_CASE
            )    ,			(
                _SCREAMING_SNAKE_CASE
            )    ,			(
                _SCREAMING_SNAKE_CASE
            )    ,			
            )       :	Optional[Any]						    =							config_and_inputs
            _SCREAMING_SNAKE_CASE       :	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 			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     UpperCamelCase_			(      self							)			->     Union[str, Any]:
            _SCREAMING_SNAKE_CASE       :	int						    =							DecisionTransformerModelTester(self							)
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							ConfigTester(self    ,						config_class=__lowerCamelCase    ,						hidden_size=3_7							)




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




       def     UpperCamelCase_			(      self							)			->     List[Any]:
            _SCREAMING_SNAKE_CASE       :	str						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*__lowerCamelCase							)




       @slow
       def     UpperCamelCase_			(      self							)			->     Dict:
            for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _SCREAMING_SNAKE_CASE       :	str						    =							DecisionTransformerModel.from_pretrained(__lowerCamelCase							)
                 self.assertIsNotNone(__lowerCamelCase							)







       def     UpperCamelCase_			(      self							)			->     Optional[int]:
            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Dict						    =							self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 _SCREAMING_SNAKE_CASE       :	Tuple						    =							model_class(__lowerCamelCase							)
                 _SCREAMING_SNAKE_CASE       :	List[str]						    =							inspect.signature(model.forward							)
                 # signature.parameters is an OrderedDict => so arg_names order is deterministic
                 _SCREAMING_SNAKE_CASE       :	int						    =							[*signature.parameters.keys()]

                 _SCREAMING_SNAKE_CASE       :	Tuple						    =							[
                     "states",
                     "actions",
                     "rewards",
                     "returns_to_go",
                     "timesteps",
                     "attention_mask",
                 ]

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





@require_torch
class 			lowerCAmelCase__( unittest.TestCase    ):






       '''simple docstring'''


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

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							torch.tensor(
                [[0.24_2793, -0.2869_3074, 0.874_2613], [0.6781_5274, -0.0810_1085, -0.1295_2147]]    ,						device=__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							torch.tensor(__lowerCamelCase    ,						device=__lowerCamelCase    ,						dtype=torch.floataa							).reshape(1    ,						1    ,						1							)
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							state
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							torch.zeros(1    ,						0    ,						config.act_dim    ,						device=__lowerCamelCase    ,						dtype=torch.floataa							)
            _SCREAMING_SNAKE_CASE       :	Dict						    =							torch.zeros(1    ,						0    ,						device=__lowerCamelCase    ,						dtype=torch.floataa							)
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							torch.tensor(0    ,						device=__lowerCamelCase    ,						dtype=torch.long							).reshape(1    ,						1							)

            for step in range(__lowerCamelCase							):
                 _SCREAMING_SNAKE_CASE       :	int						    =							torch.cat([actions, torch.zeros(1    ,						1    ,						config.act_dim    ,						device=__lowerCamelCase							)]    ,						dim=1							)
                 _SCREAMING_SNAKE_CASE       :	Any						    =							torch.cat([rewards, torch.zeros(1    ,						1    ,						device=__lowerCamelCase							)]    ,						dim=1							)

                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							torch.ones(1    ,						states.shape[1]							).to(dtype=torch.long    ,						device=states.device							)

                 with torch.no_grad():
                      _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							model(
                          states=__lowerCamelCase    ,						actions=__lowerCamelCase    ,						rewards=__lowerCamelCase    ,						returns_to_go=__lowerCamelCase    ,						timesteps=__lowerCamelCase    ,						attention_mask=__lowerCamelCase    ,						return_dict=__lowerCamelCase    ,						)

                 self.assertEqual(action_pred.shape    ,						actions.shape							)
                 self.assertTrue(torch.allclose(action_pred[0, -1]    ,						expected_outputs[step]    ,						atol=1E-4							)							)
                 _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Tuple						    =							(  # env.step(action)
                     torch.randn(1    ,						1    ,						config.state_dim							).to(device=__lowerCamelCase    ,						dtype=torch.floataa							),
                     1.0,
                     False,
                     {},
                 )

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