


'''simple docstring'''







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






      def __init__( self			:Optional[int]     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:List[Any]=1_0_0     ,						SCREAMING_SNAKE_CASE			:str=1_3     ,						SCREAMING_SNAKE_CASE			:List[str]=3_0     ,						SCREAMING_SNAKE_CASE			:str=2     ,						SCREAMING_SNAKE_CASE			:int=3     ,						SCREAMING_SNAKE_CASE			:Tuple=True     ,						SCREAMING_SNAKE_CASE			:List[str]=True     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]=3_2     ,						SCREAMING_SNAKE_CASE			:Dict=4     ,						SCREAMING_SNAKE_CASE			:Optional[Any]=4     ,						SCREAMING_SNAKE_CASE			:str=3_7     ,						SCREAMING_SNAKE_CASE			:Tuple="gelu"     ,						SCREAMING_SNAKE_CASE			:List[Any]=0.1     ,						SCREAMING_SNAKE_CASE			:Optional[int]=0.1     ,						SCREAMING_SNAKE_CASE			:List[Any]=1_0     ,						SCREAMING_SNAKE_CASE			:Any=0.02     ,						SCREAMING_SNAKE_CASE			:Any=3     ,						SCREAMING_SNAKE_CASE			:Dict=None     ,						SCREAMING_SNAKE_CASE			:List[Any]=[0, 1, 2, 3]     ,						)       ->					Dict:




             '''simple docstring'''
             _a    : Tuple     					=parent
             _a    : Dict     					=1_0_0
             _a    : List[Any]     					=batch_size
             _a    : Dict     					=image_size
             _a    : Union[str, Any]     					=patch_size
             _a    : Any     					=num_channels
             _a    : Optional[int]     					=is_training
             _a    : Optional[int]     					=use_labels
             _a    : List[str]     					=hidden_size
             _a    : Dict     					=num_hidden_layers
             _a    : Any     					=num_attention_heads
             _a    : int     					=intermediate_size
             _a    : Dict     					=hidden_act
             _a    : str     					=hidden_dropout_prob
             _a    : int     					=attention_probs_dropout_prob
             _a    : Optional[Any]     					=type_sequence_label_size
             _a    : Optional[int]     					=initializer_range
             _a    : Optional[int]     					=scope
             _a    : Tuple     					=out_indices
             _a    : Tuple     					=num_labels

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






      def __UpperCAmelCase       ( self			:Any	)       ->					List[str]:




             '''simple docstring'''
             _a    : Tuple     					=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

             _a    : Union[str, Any]     					=self.get_config()

             return config, pixel_values, labels, pixel_labels






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					List[Any]:




             '''simple docstring'''
             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=SCREAMING_SNAKE_CASE     ,						initializer_range=self.initializer_range     ,						out_indices=self.out_indices     ,						)






      def __UpperCAmelCase       ( self			:int     ,						SCREAMING_SNAKE_CASE			:int     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:int     ,						SCREAMING_SNAKE_CASE			:Dict	)       ->					Any:




             '''simple docstring'''
             _a    : List[Any]     					=BeitModel(config=SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : Optional[Any]     					=model(SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(result.last_hidden_state.shape     ,						(self.batch_size, self.seq_length, self.hidden_size)	)






      def __UpperCAmelCase       ( self			:Any     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]     ,						SCREAMING_SNAKE_CASE			:Any	)       ->					str:




             '''simple docstring'''
             _a    : Optional[Any]     					=BeitForMaskedImageModeling(config=SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : List[Any]     					=model(SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.seq_length - 1, self.vocab_size)	)






      def __UpperCAmelCase       ( self			:List[str]     ,						SCREAMING_SNAKE_CASE			:Optional[int]     ,						SCREAMING_SNAKE_CASE			:List[Any]     ,						SCREAMING_SNAKE_CASE			:List[Any]     ,						SCREAMING_SNAKE_CASE			:Tuple	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : str     					=self.type_sequence_label_size
             _a    : Union[str, Any]     					=BeitForImageClassification(SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : Union[str, Any]     					=model(SCREAMING_SNAKE_CASE     ,						labels=SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.type_sequence_label_size)	)

             # test greyscale images
             _a    : Optional[Any]     					=1
             _a    : Optional[int]     					=BeitForImageClassification(SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()

             _a    : str     					=floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
             _a    : Tuple     					=model(SCREAMING_SNAKE_CASE     ,						labels=SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.type_sequence_label_size)	)






      def __UpperCAmelCase       ( self			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:Optional[int]     ,						SCREAMING_SNAKE_CASE			:int     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]	)       ->					str:




             '''simple docstring'''
             _a    : Any     					=self.num_labels
             _a    : List[Any]     					=BeitForSemanticSegmentation(SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : List[Any]     					=model(SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(
                 result.logits.shape     ,						(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)	)
             _a    : Dict     					=model(SCREAMING_SNAKE_CASE     ,						labels=SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(
                 result.logits.shape     ,						(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)	)






      def __UpperCAmelCase       ( self			:Tuple	)       ->					List[str]:




             '''simple docstring'''
             _a    : int     					=self.prepare_config_and_inputs()
             _a					,    _a					,    _a					,    _a    : List[str]     					=config_and_inputs
             _a    : Dict     					={"""pixel_values""": pixel_values}
             return config, inputs_dict







@require_torch
class A__       (						UpperCAmelCase__ ,     UpperCAmelCase__ ,     unittest.TestCase      ):
      __UpperCamelCase       :  Tuple	       =							(
          (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
          if is_torch_available()
          else ()
      )
      __UpperCamelCase       :  List[str]	       =							(
          {
              "feature-extraction": BeitModel,
              "image-classification": BeitForImageClassification,
              "image-segmentation": BeitForSemanticSegmentation,
          }
          if is_torch_available()
          else {}
      )

      __UpperCamelCase       :  List[Any]	       =							False
      __UpperCamelCase       :  Optional[int]	       =							False
      __UpperCamelCase       :  List[Any]	       =							False






      def __UpperCAmelCase       ( self			:List[Any]	)       ->					Dict:




             '''simple docstring'''
             _a    : Dict     					=BeitModelTester(self	)
             _a    : List[str]     					=ConfigTester(self     ,						config_class=SCREAMING_SNAKE_CASE     ,						has_text_modality=SCREAMING_SNAKE_CASE     ,						hidden_size=3_7	)






      def __UpperCAmelCase       ( self			:str	)       ->					Optional[int]:




             '''simple docstring'''
             self.config_tester.run_common_tests()






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




             '''simple docstring'''
             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			:Union[str, Any]	)       ->					Dict:




             '''simple docstring'''
             pass






      def __UpperCAmelCase       ( self			:Any	)       ->					Tuple:




             '''simple docstring'''
             _a					,    _a    : Union[str, Any]     					=self.model_tester.prepare_config_and_inputs_for_common()

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






      def __UpperCAmelCase       ( self			:str	)       ->					Optional[int]:




             '''simple docstring'''
             _a					,    _a    : Union[str, Any]     					=self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    _a    : Optional[Any]     					=model_class(SCREAMING_SNAKE_CASE	)
                    _a    : int     					=inspect.signature(model.forward	)
                    # signature.parameters is an OrderedDict => so arg_names order is deterministic
                    _a    : Dict     					=[*signature.parameters.keys()]

                    _a    : str     					=["""pixel_values"""]
                    self.assertListEqual(arg_names[:1]     ,						SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Dict	)       ->					Optional[int]:




             '''simple docstring'''
             _a    : Tuple     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:List[Any]	)       ->					int:




             '''simple docstring'''
             _a    : Optional[Any]     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Dict	)       ->					Optional[int]:




             '''simple docstring'''
             _a    : List[Any]     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Union[str, Any]	)       ->					Any:




             '''simple docstring'''
             _a    : List[Any]     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_semantic_segmentation(*SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:str	)       ->					Optional[Any]:




             '''simple docstring'''
             if not self.model_tester.is_training:
                    return

             _a					,    _a    : int     					=self.model_tester.prepare_config_and_inputs_for_common()
             _a    : Optional[Any]     					=True

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

                    _a    : Tuple     					=model_class(SCREAMING_SNAKE_CASE	)
                    model.to(SCREAMING_SNAKE_CASE	)
                    model.train()
                    _a    : Optional[int]     					=self._prepare_for_class(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						return_labels=SCREAMING_SNAKE_CASE	)
                    _a    : List[Any]     					=model(**SCREAMING_SNAKE_CASE	).loss
                    loss.backward()






      def __UpperCAmelCase       ( self			:int	)       ->					Any:




             '''simple docstring'''
             _a					,    _a    : Optional[int]     					=self.model_tester.prepare_config_and_inputs_for_common()
             if not self.model_tester.is_training:
                    return

             _a    : List[str]     					=False
             _a    : int     					=True

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

                    _a    : Any     					=model_class(SCREAMING_SNAKE_CASE	)
                    model.gradient_checkpointing_enable()
                    model.to(SCREAMING_SNAKE_CASE	)
                    model.train()
                    _a    : Tuple     					=self._prepare_for_class(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						return_labels=SCREAMING_SNAKE_CASE	)
                    _a    : Dict     					=model(**SCREAMING_SNAKE_CASE	).loss
                    loss.backward()






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




             '''simple docstring'''
             _a					,    _a    : Any     					=self.model_tester.prepare_config_and_inputs_for_common()

             _a    : Union[str, Any]     					=_config_zero_init(SCREAMING_SNAKE_CASE	)
             for model_class in self.all_model_classes:
                    _a    : Dict     					=model_class(config=SCREAMING_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 __UpperCAmelCase       ( self			:str	)       ->					str:




             '''simple docstring'''
             for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a    : Union[str, Any]     					=BeitModel.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)


def 	SCREAMING_SNAKE_CASE_					(  )  ->			Union[str, Any]:
       _a    : List[Any]     					=Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""  )
       return image







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






      @cached_property
      def __UpperCAmelCase       ( self			:Dict	)       ->					str:




             '''simple docstring'''
             return BeitImageProcessor.from_pretrained("""microsoft/beit-base-patch16-224"""	) if is_vision_available() else None






      @slow
      def __UpperCAmelCase       ( self			:Union[str, Any]	)       ->					str:




             '''simple docstring'''
             _a    : Union[str, Any]     					=BeitForMaskedImageModeling.from_pretrained("""microsoft/beit-base-patch16-224-pt22k"""	).to(SCREAMING_SNAKE_CASE	)

             _a    : Tuple     					=self.default_image_processor
             _a    : Union[str, Any]     					=prepare_img()
             _a    : Tuple     					=image_processor(images=SCREAMING_SNAKE_CASE     ,						return_tensors="""pt"""	).pixel_values.to(SCREAMING_SNAKE_CASE	)

             # prepare bool_masked_pos
             _a    : int     					=torch.ones((1, 1_9_6)     ,						dtype=torch.bool	).to(SCREAMING_SNAKE_CASE	)

             # forward pass
             with torch.no_grad():
                    _a    : Any     					=model(pixel_values=SCREAMING_SNAKE_CASE     ,						bool_masked_pos=SCREAMING_SNAKE_CASE	)
             _a    : int     					=outputs.logits

             # verify the logits
             _a    : Any     					=torch.Size((1, 1_9_6, 8_1_9_2)	)
             self.assertEqual(logits.shape     ,						SCREAMING_SNAKE_CASE	)

             _a    : Optional[int]     					=torch.tensor(
                 [[-3.2_437, 0.5_072, -13.9_174], [-3.2_456, 0.4_948, -13.9_401], [-3.2_033, 0.5_121, -13.8_550]]	).to(SCREAMING_SNAKE_CASE	)

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






      @slow
      def __UpperCAmelCase       ( self			:int	)       ->					Tuple:




             '''simple docstring'''
             _a    : List[str]     					=BeitForImageClassification.from_pretrained("""microsoft/beit-base-patch16-224"""	).to(SCREAMING_SNAKE_CASE	)

             _a    : Dict     					=self.default_image_processor
             _a    : Tuple     					=prepare_img()
             _a    : str     					=image_processor(images=SCREAMING_SNAKE_CASE     ,						return_tensors="""pt"""	).to(SCREAMING_SNAKE_CASE	)

             # forward pass
             with torch.no_grad():
                    _a    : Union[str, Any]     					=model(**SCREAMING_SNAKE_CASE	)
             _a    : int     					=outputs.logits

             # verify the logits
             _a    : Optional[Any]     					=torch.Size((1, 1_0_0_0)	)
             self.assertEqual(logits.shape     ,						SCREAMING_SNAKE_CASE	)

             _a    : Optional[int]     					=torch.tensor([-1.2_385, -1.0_987, -1.0_108]	).to(SCREAMING_SNAKE_CASE	)

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

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






      @slow
      def __UpperCAmelCase       ( self			:Any	)       ->					Any:




             '''simple docstring'''
             _a    : int     					=BeitForImageClassification.from_pretrained("""microsoft/beit-large-patch16-224-pt22k-ft22k"""	).to(
                 SCREAMING_SNAKE_CASE	)

             _a    : List[Any]     					=self.default_image_processor
             _a    : Optional[Any]     					=prepare_img()
             _a    : Optional[Any]     					=image_processor(images=SCREAMING_SNAKE_CASE     ,						return_tensors="""pt"""	).to(SCREAMING_SNAKE_CASE	)

             # forward pass
             with torch.no_grad():
                    _a    : Any     					=model(**SCREAMING_SNAKE_CASE	)
             _a    : Union[str, Any]     					=outputs.logits

             # verify the logits
             _a    : Tuple     					=torch.Size((1, 2_1_8_4_1)	)
             self.assertEqual(logits.shape     ,						SCREAMING_SNAKE_CASE	)

             _a    : List[str]     					=torch.tensor([1.6_881, -0.2_787, 0.5_901]	).to(SCREAMING_SNAKE_CASE	)

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

             _a    : List[str]     					=2_3_9_6
             self.assertEqual(logits.argmax(-1	).item()     ,						SCREAMING_SNAKE_CASE	)






      @slow
      def __UpperCAmelCase       ( self			:int	)       ->					Any:




             '''simple docstring'''
             _a    : Union[str, Any]     					=BeitForSemanticSegmentation.from_pretrained("""microsoft/beit-base-finetuned-ade-640-640"""	)
             _a    : List[Any]     					=model.to(SCREAMING_SNAKE_CASE	)

             _a    : Tuple     					=BeitImageProcessor(do_resize=SCREAMING_SNAKE_CASE     ,						size=6_4_0     ,						do_center_crop=SCREAMING_SNAKE_CASE	)

             _a    : Tuple     					=load_dataset("""hf-internal-testing/fixtures_ade20k"""     ,						split="""test"""	)
             _a    : Tuple     					=Image.open(ds[0]["""file"""]	)
             _a    : List[str]     					=image_processor(images=SCREAMING_SNAKE_CASE     ,						return_tensors="""pt"""	).to(SCREAMING_SNAKE_CASE	)

             # forward pass
             with torch.no_grad():
                    _a    : Any     					=model(**SCREAMING_SNAKE_CASE	)
             _a    : Tuple     					=outputs.logits

             # verify the logits
             _a    : List[str]     					=torch.Size((1, 1_5_0, 1_6_0, 1_6_0)	)
             self.assertEqual(logits.shape     ,						SCREAMING_SNAKE_CASE	)

             _a    : Optional[Any]     					=version.parse(PIL.__version__	) < version.parse("""9.0.0"""	)

             if is_pillow_less_than_a:
                    _a    : Dict     					=torch.tensor(
                        [
                            [[-4.9_225, -2.3_954, -3.0_522], [-2.8_822, -1.0_046, -1.7_561], [-2.9_549, -1.3_228, -2.1_347]],
                            [[-5.8_168, -3.4_129, -4.0_778], [-3.8_651, -2.2_214, -3.0_277], [-3.8_356, -2.4_643, -3.3_535]],
                            [[-0.0_078, 3.9_952, 4.0_754], [2.9_856, 4.6_944, 5.0_035], [3.2_413, 4.7_813, 4.9_969]],
                        ]     ,						device=SCREAMING_SNAKE_CASE     ,						)
             else:
                    _a    : List[str]     					=torch.tensor(
                        [
                            [[-4.8_960, -2.3_688, -3.0_355], [-2.8_478, -0.9_836, -1.7_418], [-2.9_449, -1.3_332, -2.1_456]],
                            [[-5.8_081, -3.4_124, -4.1_006], [-3.8_561, -2.2_081, -3.0_323], [-3.8_365, -2.4_601, -3.3_669]],
                            [[-0.0_309, 3.9_868, 4.0_540], [2.9_640, 4.6_877, 4.9_976], [3.2_081, 4.7_690, 4.9_942]],
                        ]     ,						device=SCREAMING_SNAKE_CASE     ,						)

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






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




             '''simple docstring'''
             _a    : Optional[int]     					=BeitForSemanticSegmentation.from_pretrained("""microsoft/beit-base-finetuned-ade-640-640"""	)
             _a    : List[str]     					=model.to(SCREAMING_SNAKE_CASE	)

             _a    : List[str]     					=BeitImageProcessor(do_resize=SCREAMING_SNAKE_CASE     ,						size=6_4_0     ,						do_center_crop=SCREAMING_SNAKE_CASE	)

             _a    : List[str]     					=load_dataset("""hf-internal-testing/fixtures_ade20k"""     ,						split="""test"""	)
             _a    : Tuple     					=Image.open(ds[0]["""file"""]	)
             _a    : Tuple     					=image_processor(images=SCREAMING_SNAKE_CASE     ,						return_tensors="""pt"""	).to(SCREAMING_SNAKE_CASE	)

             # forward pass
             with torch.no_grad():
                    _a    : Dict     					=model(**SCREAMING_SNAKE_CASE	)

             _a    : Dict     					=outputs.logits.detach().cpu()

             _a    : Optional[Any]     					=image_processor.post_process_semantic_segmentation(outputs=SCREAMING_SNAKE_CASE     ,						target_sizes=[(5_0_0, 3_0_0)]	)
             _a    : List[str]     					=torch.Size((5_0_0, 3_0_0)	)
             self.assertEqual(segmentation[0].shape     ,						SCREAMING_SNAKE_CASE	)

             _a    : List[str]     					=image_processor.post_process_semantic_segmentation(outputs=SCREAMING_SNAKE_CASE	)
             _a    : Any     					=torch.Size((1_6_0, 1_6_0)	)
             self.assertEqual(segmentation[0].shape     ,						SCREAMING_SNAKE_CASE	)






'''simple docstring'''







class A__       :






      def __init__( self			:List[Any]	)       ->					None:




             '''simple docstring'''
             _a    : dict[str, TrieNode]     					={}  # Mapping from char to TrieNode
             _a    : List[str]     					=False






      def __UpperCAmelCase       ( self			:Tuple     ,						SCREAMING_SNAKE_CASE			:list[str]	)       ->					None:




             '''simple docstring'''
             for word in words:
                    self.insert(SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Union[str, Any]     ,						SCREAMING_SNAKE_CASE			:str	)       ->					None:




             '''simple docstring'''
             _a    : str     					=self
             for char in word:
                    if char not in curr.nodes:
                           _a    : Dict     					=TrieNode()
                    _a    : List[Any]     					=curr.nodes[char]
             _a    : int     					=True






      def __UpperCAmelCase       ( self			:Union[str, Any]     ,						SCREAMING_SNAKE_CASE			:str	)       ->					bool:




             '''simple docstring'''
             _a    : int     					=self
             for char in word:
                    if char not in curr.nodes:
                           return False
                    _a    : List[Any]     					=curr.nodes[char]
             return curr.is_leaf






      def __UpperCAmelCase       ( self			:Dict     ,						SCREAMING_SNAKE_CASE			:str	)       ->					None:




             '''simple docstring'''

             def _delete(SCREAMING_SNAKE_CASE			:TrieNode     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:int	) -> bool:
                    if index == len(SCREAMING_SNAKE_CASE	):
                           # If word does not exist
                           if not curr.is_leaf:
                                  return False
                           _a    : Any     					=False
                           return len(curr.nodes	) == 0
                    _a    : int     					=word[index]
                    _a    : int     					=curr.nodes.get(SCREAMING_SNAKE_CASE	)
                    # If char not in current trie node
                    if not char_node:
                           return False
                    # Flag to check if node can be deleted
                    _a    : List[Any]     					=_delete(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE     ,						index + 1	)
                    if delete_curr:
                           del curr.nodes[char]
                           return len(curr.nodes	) == 0
                    return delete_curr

             _delete(self     ,						SCREAMING_SNAKE_CASE     ,						0	)


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	TrieNode    ,_UpperCAmelCase	:	str  )  ->			None:
       if node.is_leaf:
              print(_UpperCAmelCase    ,end=""" """  )

       for key, value in node.nodes.items():
              print_words(_UpperCAmelCase    ,word + key  )


def 	SCREAMING_SNAKE_CASE_					(  )  ->			bool:
       _a    : List[str]     					="""banana bananas bandana band apple all beast""".split()
       _a    : List[Any]     					=TrieNode()
       root.insert_many(_UpperCAmelCase  )
       # print_words(root, "")
       assert all(root.find(_UpperCAmelCase  ) for word in words  )
       assert root.find("""banana"""  )
       assert not root.find("""bandanas"""  )
       assert not root.find("""apps"""  )
       assert root.find("""apple"""  )
       assert root.find("""all"""  )
       root.delete("""all"""  )
       assert not root.find("""all"""  )
       root.delete("""banana"""  )
       assert not root.find("""banana"""  )
       assert root.find("""bananas"""  )
       return True


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	str    ,_UpperCAmelCase	:	bool  )  ->			None:
       print(str(_UpperCAmelCase  )    ,"""works!""" if passes else """doesn't work :("""  )


def 	SCREAMING_SNAKE_CASE_					(  )  ->			None:
       assert test_trie()


def 	SCREAMING_SNAKE_CASE_					(  )  ->			None:
       print_results("""Testing trie functionality"""    ,test_trie()  )


if __name__ == "__main__":
   main()


