

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 				lowercase       :




				def __init__(      self   ,     snake_case   ,     snake_case=100   ,     snake_case=13   ,     snake_case=30   ,     snake_case=2   ,     snake_case=3   ,     snake_case=True   ,     snake_case=True   ,     snake_case=32   ,     snake_case=4   ,     snake_case=4   ,     snake_case=37   ,     snake_case="gelu"   ,     snake_case=0.1   ,     snake_case=0.1   ,     snake_case=10   ,     snake_case=0.02   ,     snake_case=3   ,     snake_case=None   ,     snake_case=[0, 1, 2, 3]   ,     ):
											snake_case_           =					parent
											snake_case_           =					100
											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_           =					scope
											snake_case_           =					out_indices
											snake_case_           =					num_labels

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




				def 					a							(      self ):
											snake_case_           =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

											snake_case_           =					None
											snake_case_           =					None
											if self.use_labels:
																		snake_case_           =					ids_tensor([self.batch_size]   ,     self.type_sequence_label_size )
																		snake_case_           =					ids_tensor([self.batch_size, self.image_size, self.image_size]   ,     self.num_labels )

											snake_case_           =					self.get_config()

											return config, pixel_values, labels, pixel_labels




				def 					a							(      self ):
											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=snake_case   ,     initializer_range=self.initializer_range   ,     out_indices=self.out_indices   ,     )




				def 					a							(      self   ,     snake_case   ,     snake_case   ,     snake_case   ,     snake_case ):
											snake_case_           =					BeitModel(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.seq_length, self.hidden_size) )




				def 					a							(      self   ,     snake_case   ,     snake_case   ,     snake_case   ,     snake_case ):
											snake_case_           =					BeitForMaskedImageModeling(config=snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					model(snake_case )
											self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length - 1, self.vocab_size) )




				def 					a							(      self   ,     snake_case   ,     snake_case   ,     snake_case   ,     snake_case ):
											snake_case_           =					self.type_sequence_label_size
											snake_case_           =					BeitForImageClassification(snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					model(snake_case   ,     labels=snake_case )
											self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.type_sequence_label_size) )

											# test greyscale images
											snake_case_           =					1
											snake_case_           =					BeitForImageClassification(snake_case )
											model.to(snake_case )
											model.eval()

											snake_case_           =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
											snake_case_           =					model(snake_case   ,     labels=snake_case )
											self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.type_sequence_label_size) )




				def 					a							(      self   ,     snake_case   ,     snake_case   ,     snake_case   ,     snake_case ):
											snake_case_           =					self.num_labels
											snake_case_           =					BeitForSemanticSegmentation(snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					model(snake_case )
											self.parent.assertEqual(
											    result.logits.shape   ,     (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2) )
											snake_case_           =					model(snake_case   ,     labels=snake_case )
											self.parent.assertEqual(
											    result.logits.shape   ,     (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2) )




				def 					a							(      self ):
											snake_case_           =					self.prepare_config_and_inputs()
											snake_case_       ,					snake_case_       ,					snake_case_       ,					snake_case_           =					config_and_inputs
											snake_case_           =					{'pixel_values': pixel_values}
											return config, inputs_dict







@require_torch
class 				lowercase       ( lowercase_			,						lowercase_			,						unittest.TestCase    ):
				__SCREAMING_SNAKE_CASE      :				str     =    (
				    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)
				__SCREAMING_SNAKE_CASE      :				List[Any]     =    (
				    {
				        '''feature-extraction''': BeitModel,
				        '''image-classification''': BeitForImageClassification,
				        '''image-segmentation''': BeitForSemanticSegmentation,
				    }
				    if is_torch_available()
				    else {}
				)

				__SCREAMING_SNAKE_CASE      :				List[str]     =    False
				__SCREAMING_SNAKE_CASE      :				Dict     =    False
				__SCREAMING_SNAKE_CASE      :				str     =    False




				def 					a							(      self ):
											snake_case_           =					BeitModelTester(self )
											snake_case_           =					ConfigTester(self   ,     config_class=snake_case   ,     has_text_modality=snake_case   ,     hidden_size=37 )




				def 					a							(      self ):
											self.config_tester.run_common_tests()




				@unittest.skip(reason='BEiT does not use inputs_embeds' )
				def 					a							(      self ):
											pass




				@require_torch_multi_gpu
				@unittest.skip(reason='BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`' )
				def 					a							(      self ):
											pass




				def 					a							(      self ):
											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 					a							(      self ):
											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 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_semantic_segmentation(*snake_case )




				def 					a							(      self ):
											if not self.model_tester.is_training:
																		return

											snake_case_       ,					snake_case_           =					self.model_tester.prepare_config_and_inputs_for_common()
											snake_case_           =					True

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

																		snake_case_           =					model_class(snake_case )
																		model.to(snake_case )
																		model.train()
																		snake_case_           =					self._prepare_for_class(snake_case   ,     snake_case   ,     return_labels=snake_case )
																		snake_case_           =					model(**snake_case ).loss
																		loss.backward()




				def 					a							(      self ):
											snake_case_       ,					snake_case_           =					self.model_tester.prepare_config_and_inputs_for_common()
											if not self.model_tester.is_training:
																		return

											snake_case_           =					False
											snake_case_           =					True

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

																		snake_case_           =					model_class(snake_case )
																		model.gradient_checkpointing_enable()
																		model.to(snake_case )
																		model.train()
																		snake_case_           =					self._prepare_for_class(snake_case   ,     snake_case   ,     return_labels=snake_case )
																		snake_case_           =					model(**snake_case ).loss
																		loss.backward()




				def 					a							(      self ):
											snake_case_       ,					snake_case_           =					self.model_tester.prepare_config_and_inputs_for_common()

											snake_case_           =					_config_zero_init(snake_case )
											for model_class in self.all_model_classes:
																		snake_case_           =					model_class(config=snake_case )
																		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 					a							(      self ):
											for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		snake_case_           =					BeitModel.from_pretrained(snake_case )
																		self.assertIsNotNone(snake_case )





def       __lowerCamelCase			(				):







							'''simple docstring'''






							snake_case_           =					Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
							return image







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




				@cached_property
				def 					a							(      self ):
											return BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224' ) if is_vision_available() else None




				@slow
				def 					a							(      self ):
											snake_case_           =					BeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k' ).to(snake_case )

											snake_case_           =					self.default_image_processor
											snake_case_           =					prepare_img()
											snake_case_           =					image_processor(images=snake_case   ,     return_tensors='pt' ).pixel_values.to(snake_case )

											# prepare bool_masked_pos
											snake_case_           =					torch.ones((1, 196)   ,     dtype=torch.bool ).to(snake_case )

											# forward pass
											with torch.no_grad():
																		snake_case_           =					model(pixel_values=snake_case   ,     bool_masked_pos=snake_case )
											snake_case_           =					outputs.logits

											# verify the logits
											snake_case_           =					torch.Size((1, 196, 8192) )
											self.assertEqual(logits.shape   ,     snake_case )

											snake_case_           =					torch.tensor(
											    [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]] ).to(snake_case )

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




				@slow
				def 					a							(      self ):
											snake_case_           =					BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224' ).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(**snake_case )
											snake_case_           =					outputs.logits

											# verify the logits
											snake_case_           =					torch.Size((1, 1000) )
											self.assertEqual(logits.shape   ,     snake_case )

											snake_case_           =					torch.tensor([-1.23_85, -1.09_87, -1.01_08] ).to(snake_case )

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

											snake_case_           =					281
											self.assertEqual(logits.argmax(-1 ).item()   ,     snake_case )




				@slow
				def 					a							(      self ):
											snake_case_           =					BeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k' ).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(**snake_case )
											snake_case_           =					outputs.logits

											# verify the logits
											snake_case_           =					torch.Size((1, 2_1841) )
											self.assertEqual(logits.shape   ,     snake_case )

											snake_case_           =					torch.tensor([1.68_81, -0.27_87, 0.59_01] ).to(snake_case )

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

											snake_case_           =					2396
											self.assertEqual(logits.argmax(-1 ).item()   ,     snake_case )




				@slow
				def 					a							(      self ):
											snake_case_           =					BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640' )
											snake_case_           =					model.to(snake_case )

											snake_case_           =					BeitImageProcessor(do_resize=snake_case   ,     size=640   ,     do_center_crop=snake_case )

											snake_case_           =					load_dataset('hf-internal-testing/fixtures_ade20k'   ,     split='test' )
											snake_case_           =					Image.open(ds[0]['file'] )
											snake_case_           =					image_processor(images=snake_case   ,     return_tensors='pt' ).to(snake_case )

											# forward pass
											with torch.no_grad():
																		snake_case_           =					model(**snake_case )
											snake_case_           =					outputs.logits

											# verify the logits
											snake_case_           =					torch.Size((1, 150, 160, 160) )
											self.assertEqual(logits.shape   ,     snake_case )

											snake_case_           =					version.parse(PIL.__version__ ) < version.parse('9.0.0' )

											if is_pillow_less_than_a:
																		snake_case_           =					torch.tensor(
																		    [
																		        [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
																		        [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
																		        [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
																		    ]   ,     device=snake_case   ,     )
											else:
																		snake_case_           =					torch.tensor(
																		    [
																		        [[-4.89_60, -2.36_88, -3.03_55], [-2.84_78, -0.98_36, -1.74_18], [-2.94_49, -1.33_32, -2.14_56]],
																		        [[-5.80_81, -3.41_24, -4.10_06], [-3.85_61, -2.20_81, -3.03_23], [-3.83_65, -2.46_01, -3.36_69]],
																		        [[-0.03_09, 3.98_68, 4.05_40], [2.96_40, 4.68_77, 4.99_76], [3.20_81, 4.76_90, 4.99_42]],
																		    ]   ,     device=snake_case   ,     )

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




				@slow
				def 					a							(      self ):
											snake_case_           =					BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640' )
											snake_case_           =					model.to(snake_case )

											snake_case_           =					BeitImageProcessor(do_resize=snake_case   ,     size=640   ,     do_center_crop=snake_case )

											snake_case_           =					load_dataset('hf-internal-testing/fixtures_ade20k'   ,     split='test' )
											snake_case_           =					Image.open(ds[0]['file'] )
											snake_case_           =					image_processor(images=snake_case   ,     return_tensors='pt' ).to(snake_case )

											# forward pass
											with torch.no_grad():
																		snake_case_           =					model(**snake_case )

											snake_case_           =					outputs.logits.detach().cpu()

											snake_case_           =					image_processor.post_process_semantic_segmentation(outputs=snake_case   ,     target_sizes=[(500, 300)] )
											snake_case_           =					torch.Size((500, 300) )
											self.assertEqual(segmentation[0].shape   ,     snake_case )

											snake_case_           =					image_processor.post_process_semantic_segmentation(outputs=snake_case )
											snake_case_           =					torch.Size((160, 160) )
											self.assertEqual(segmentation[0].shape   ,     snake_case )







from unittest import TestCase

from datasets import Sequence, Value
from datasets.arrow_dataset import Dataset







class 				lowercase       ( lowercase_    ):




				def 					a							(      self ):
											return [
											    {"col_1": 3, "col_2": "a"},
											    {"col_1": 2, "col_2": "b"},
											    {"col_1": 1, "col_2": "c"},
											    {"col_1": 0, "col_2": "d"},
											]




				def 					a							(      self ):
											snake_case_           =					{'col_1': [3, 2, 1, 0], 'col_2': ['a', 'b', 'c', 'd']}
											return Dataset.from_dict(snake_case )




				def 					a							(      self ):
											snake_case_           =					self._create_example_records()
											snake_case_           =					Dataset.from_list(snake_case )
											self.assertListEqual(dset.column_names   ,     ['col_1', 'col_2'] )
											for i, r in enumerate(snake_case ):
																		self.assertDictEqual(snake_case   ,     example_records[i] )




				def 					a							(      self ):
											snake_case_           =					self._create_example_records()
											snake_case_           =					Dataset.from_list(snake_case )
											snake_case_           =					Dataset.from_dict({k: [r[k] for r in example_records] for k in example_records[0]} )
											self.assertEqual(dset.info   ,     dset_from_dict.info )




				def 					a							(      self ):  # checks what happens with missing columns
											snake_case_           =					[{'col_1': 1}, {'col_2': 'x'}]
											snake_case_           =					Dataset.from_list(snake_case )
											self.assertDictEqual(dset[0]   ,     {'col_1': 1} )
											self.assertDictEqual(dset[1]   ,     {'col_1': None} )  # NB: first record is used for columns




				def 					a							(      self ):  # checks if the type can be inferred from the second record
											snake_case_           =					[{'col_1': []}, {'col_1': [1, 2]}]
											snake_case_           =					Dataset.from_list(snake_case )
											self.assertEqual(dset.info.features['col_1']   ,     Sequence(Value('int64' ) ) )




				def 					a							(      self ):
											snake_case_           =					Dataset.from_list([] )
											self.assertEqual(len(snake_case )   ,     0 )
											self.assertListEqual(dset.column_names   ,     [] )




