





"""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    _SCREAMING_SNAKE_CASE:




       def __init__(		self						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__=1_00						,SCREAMING_SNAKE_CASE__=13						,SCREAMING_SNAKE_CASE__=30						,SCREAMING_SNAKE_CASE__=2						,SCREAMING_SNAKE_CASE__=3						,SCREAMING_SNAKE_CASE__=True						,SCREAMING_SNAKE_CASE__=True						,SCREAMING_SNAKE_CASE__=32						,SCREAMING_SNAKE_CASE__=4						,SCREAMING_SNAKE_CASE__=4						,SCREAMING_SNAKE_CASE__=37						,SCREAMING_SNAKE_CASE__="gelu"						,SCREAMING_SNAKE_CASE__=0.1						,SCREAMING_SNAKE_CASE__=0.1						,SCREAMING_SNAKE_CASE__=10						,SCREAMING_SNAKE_CASE__=0.0_2						,SCREAMING_SNAKE_CASE__=3						,SCREAMING_SNAKE_CASE__=None						,SCREAMING_SNAKE_CASE__=[0, 1, 2, 3]						,)     -> Tuple:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Tuple													=    parent
             __SCREAMING_SNAKE_CASE  :List[Any]													=    1_00
             __SCREAMING_SNAKE_CASE  :Dict													=    batch_size
             __SCREAMING_SNAKE_CASE  :Optional[Any]													=    image_size
             __SCREAMING_SNAKE_CASE  :Optional[Any]													=    patch_size
             __SCREAMING_SNAKE_CASE  :Tuple													=    num_channels
             __SCREAMING_SNAKE_CASE  :Dict													=    is_training
             __SCREAMING_SNAKE_CASE  :Optional[Any]													=    use_labels
             __SCREAMING_SNAKE_CASE  :Dict													=    hidden_size
             __SCREAMING_SNAKE_CASE  :Tuple													=    num_hidden_layers
             __SCREAMING_SNAKE_CASE  :List[str]													=    num_attention_heads
             __SCREAMING_SNAKE_CASE  :Dict													=    intermediate_size
             __SCREAMING_SNAKE_CASE  :str													=    hidden_act
             __SCREAMING_SNAKE_CASE  :Any													=    hidden_dropout_prob
             __SCREAMING_SNAKE_CASE  :Tuple													=    attention_probs_dropout_prob
             __SCREAMING_SNAKE_CASE  :int													=    type_sequence_label_size
             __SCREAMING_SNAKE_CASE  :str													=    initializer_range
             __SCREAMING_SNAKE_CASE  :Dict													=    scope
             __SCREAMING_SNAKE_CASE  :Dict													=    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  :Any													=    (image_size // patch_size) ** 2
             __SCREAMING_SNAKE_CASE  :Dict													=    num_patches + 1




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :int													=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

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

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

             return config, pixel_values, labels, pixel_labels




       def        _UpperCamelCase	(		self				)     -> 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=A_						,initializer_range=self.initializer_range						,out_indices=self.out_indices						,)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :int													=    BeitModel(config=A_				)
             model.to(A_				)
             model.eval()
             __SCREAMING_SNAKE_CASE  :List[str]													=    model(A_				)
             self.parent.assertEqual(result.last_hidden_state.shape						,(self.batch_size, self.seq_length, self.hidden_size)				)




       def        _UpperCamelCase	(		self						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__				)     -> str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Optional[int]													=    BeitForMaskedImageModeling(config=A_				)
             model.to(A_				)
             model.eval()
             __SCREAMING_SNAKE_CASE  :int													=    model(A_				)
             self.parent.assertEqual(result.logits.shape						,(self.batch_size, self.seq_length - 1, self.vocab_size)				)




       def        _UpperCamelCase	(		self						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__				)     -> Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Optional[int]													=    self.type_sequence_label_size
             __SCREAMING_SNAKE_CASE  :str													=    BeitForImageClassification(A_				)
             model.to(A_				)
             model.eval()
             __SCREAMING_SNAKE_CASE  :int													=    model(A_						,labels=A_				)
             self.parent.assertEqual(result.logits.shape						,(self.batch_size, self.type_sequence_label_size)				)

             # test greyscale images
             __SCREAMING_SNAKE_CASE  :str													=    1
             __SCREAMING_SNAKE_CASE  :List[str]													=    BeitForImageClassification(A_				)
             model.to(A_				)
             model.eval()

             __SCREAMING_SNAKE_CASE  :Optional[Any]													=    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
             __SCREAMING_SNAKE_CASE  :Dict													=    model(A_						,labels=A_				)
             self.parent.assertEqual(result.logits.shape						,(self.batch_size, self.type_sequence_label_size)				)




       def        _UpperCamelCase	(		self						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__						,SCREAMING_SNAKE_CASE__				)     -> Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Dict													=    self.num_labels
             __SCREAMING_SNAKE_CASE  :Any													=    BeitForSemanticSegmentation(A_				)
             model.to(A_				)
             model.eval()
             __SCREAMING_SNAKE_CASE  :str													=    model(A_				)
             self.parent.assertEqual(
                 result.logits.shape						,(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)				)
             __SCREAMING_SNAKE_CASE  :Any													=    model(A_						,labels=A_				)
             self.parent.assertEqual(
                 result.logits.shape						,(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)				)




       def        _UpperCamelCase	(		self				)     -> int:





             """simple docstring"""






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






@require_torch
class    _SCREAMING_SNAKE_CASE(  _a       ,      _a       ,      unittest.TestCase    ):
       SCREAMING_SNAKE_CASE_  : List[Any]      		=    (
           (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
           if is_torch_available()
           else ()
       )
       SCREAMING_SNAKE_CASE_  : Dict      		=    (
           {
               """feature-extraction""": BeitModel,
               """image-classification""": BeitForImageClassification,
               """image-segmentation""": BeitForSemanticSegmentation,
           }
           if is_torch_available()
           else {}
       )

       SCREAMING_SNAKE_CASE_  : Tuple      		=    False
       SCREAMING_SNAKE_CASE_  : List[str]      		=    False
       SCREAMING_SNAKE_CASE_  : int      		=    False




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :int													=    BeitModelTester(self				)
             __SCREAMING_SNAKE_CASE  :List[str]													=    ConfigTester(self						,config_class=A_						,has_text_modality=A_						,hidden_size=37				)




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





             """simple docstring"""






             self.config_tester.run_common_tests()




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





             """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				)     -> int:





             """simple docstring"""






             pass




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





             """simple docstring"""






             __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  :Union[str, Any]													=    model_class(A_				)
                   self.assertIsInstance(model.get_input_embeddings()						,(nn.Module)				)
                   __SCREAMING_SNAKE_CASE  :Optional[int]													=    model.get_output_embeddings()
                   self.assertTrue(x is None or isinstance(A_						,nn.Linear				)				)




       def        _UpperCamelCase	(		self				)     -> str:





             """simple docstring"""






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

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

                   __SCREAMING_SNAKE_CASE  :int													=    ['''pixel_values''']
                   self.assertListEqual(arg_names[:1]						,A_				)




       def        _UpperCamelCase	(		self				)     -> int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Optional[Any]													=    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*A_				)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Dict													=    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*A_				)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Union[str, Any]													=    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*A_				)




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :int													=    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_semantic_segmentation(*A_				)




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





             """simple docstring"""






             if not self.model_tester.is_training:
                   return

             __SCREAMING_SNAKE_CASE							,__SCREAMING_SNAKE_CASE  :Dict													=    self.model_tester.prepare_config_and_inputs_for_common()
             __SCREAMING_SNAKE_CASE  :str													=    True

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

                   __SCREAMING_SNAKE_CASE  :Tuple													=    model_class(A_				)
                   model.to(A_				)
                   model.train()
                   __SCREAMING_SNAKE_CASE  :Tuple													=    self._prepare_for_class(A_						,A_						,return_labels=A_				)
                   __SCREAMING_SNAKE_CASE  :Tuple													=    model(**A_				).loss
                   loss.backward()




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE							,__SCREAMING_SNAKE_CASE  :List[Any]													=    self.model_tester.prepare_config_and_inputs_for_common()
             if not self.model_tester.is_training:
                   return

             __SCREAMING_SNAKE_CASE  :str													=    False
             __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(A_				), BeitForMaskedImageModeling]
                       or not model_class.supports_gradient_checkpointing
                   ):
                         continue

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




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE							,__SCREAMING_SNAKE_CASE  :Optional[int]													=    self.model_tester.prepare_config_and_inputs_for_common()

             __SCREAMING_SNAKE_CASE  :int													=    _config_zero_init(A_				)
             for model_class in self.all_model_classes:
                   __SCREAMING_SNAKE_CASE  :int													=    model_class(config=A_				)
                   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				)     -> Union[str, Any]:





             """simple docstring"""






             for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   __SCREAMING_SNAKE_CASE  :List[str]													=    BeitModel.from_pretrained(A_				)
                   self.assertIsNotNone(A_				)




def 	__lowerCamelCase  (				)     -> List[str]:
      __SCREAMING_SNAKE_CASE  :List[str]													=    Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
      return image







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




       @cached_property
       def        _UpperCamelCase	(		self				)     -> Any:





             """simple docstring"""






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




       @slow
       def        _UpperCamelCase	(		self				)     -> Tuple:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Optional[Any]													=    BeitForMaskedImageModeling.from_pretrained('''microsoft/beit-base-patch16-224-pt22k'''				).to(A_				)

             __SCREAMING_SNAKE_CASE  :Dict													=    self.default_image_processor
             __SCREAMING_SNAKE_CASE  :List[Any]													=    prepare_img()
             __SCREAMING_SNAKE_CASE  :Any													=    image_processor(images=A_						,return_tensors='''pt'''				).pixel_values.to(A_				)

             # prepare bool_masked_pos
             __SCREAMING_SNAKE_CASE  :Optional[int]													=    torch.ones((1, 1_96)						,dtype=torch.bool				).to(A_				)

             # forward pass
             with torch.no_grad():
                   __SCREAMING_SNAKE_CASE  :Optional[int]													=    model(pixel_values=A_						,bool_masked_pos=A_				)
             __SCREAMING_SNAKE_CASE  :Union[str, Any]													=    outputs.logits

             # verify the logits
             __SCREAMING_SNAKE_CASE  :List[str]													=    torch.Size((1, 1_96, 81_92)				)
             self.assertEqual(logits.shape						,A_				)

             __SCREAMING_SNAKE_CASE  :List[str]													=    torch.tensor(
                 [[-3.2_4_3_7, 0.5_0_7_2, -1_3.9_1_7_4], [-3.2_4_5_6, 0.4_9_4_8, -1_3.9_4_0_1], [-3.2_0_3_3, 0.5_1_2_1, -1_3.8_5_5_0]]				).to(A_				)

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




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :List[Any]													=    BeitForImageClassification.from_pretrained('''microsoft/beit-base-patch16-224'''				).to(A_				)

             __SCREAMING_SNAKE_CASE  :Tuple													=    self.default_image_processor
             __SCREAMING_SNAKE_CASE  :int													=    prepare_img()
             __SCREAMING_SNAKE_CASE  :Any													=    image_processor(images=A_						,return_tensors='''pt'''				).to(A_				)

             # forward pass
             with torch.no_grad():
                   __SCREAMING_SNAKE_CASE  :Dict													=    model(**A_				)
             __SCREAMING_SNAKE_CASE  :str													=    outputs.logits

             # verify the logits
             __SCREAMING_SNAKE_CASE  :Optional[int]													=    torch.Size((1, 10_00)				)
             self.assertEqual(logits.shape						,A_				)

             __SCREAMING_SNAKE_CASE  :Any													=    torch.tensor([-1.2_3_8_5, -1.0_9_8_7, -1.0_1_0_8]				).to(A_				)

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

             __SCREAMING_SNAKE_CASE  :List[Any]													=    2_81
             self.assertEqual(logits.argmax(-1				).item()						,A_				)




       @slow
       def        _UpperCamelCase	(		self				)     -> Tuple:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Union[str, Any]													=    BeitForImageClassification.from_pretrained('''microsoft/beit-large-patch16-224-pt22k-ft22k'''				).to(
                 A_				)

             __SCREAMING_SNAKE_CASE  :str													=    self.default_image_processor
             __SCREAMING_SNAKE_CASE  :Any													=    prepare_img()
             __SCREAMING_SNAKE_CASE  :Dict													=    image_processor(images=A_						,return_tensors='''pt'''				).to(A_				)

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

             # verify the logits
             __SCREAMING_SNAKE_CASE  :str													=    torch.Size((1, 2_18_41)				)
             self.assertEqual(logits.shape						,A_				)

             __SCREAMING_SNAKE_CASE  :Optional[int]													=    torch.tensor([1.6_8_8_1, -0.2_7_8_7, 0.5_9_0_1]				).to(A_				)

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

             __SCREAMING_SNAKE_CASE  :Tuple													=    23_96
             self.assertEqual(logits.argmax(-1				).item()						,A_				)




       @slow
       def        _UpperCamelCase	(		self				)     -> Tuple:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :Optional[Any]													=    BeitForSemanticSegmentation.from_pretrained('''microsoft/beit-base-finetuned-ade-640-640'''				)
             __SCREAMING_SNAKE_CASE  :Any													=    model.to(A_				)

             __SCREAMING_SNAKE_CASE  :int													=    BeitImageProcessor(do_resize=A_						,size=6_40						,do_center_crop=A_				)

             __SCREAMING_SNAKE_CASE  :Union[str, Any]													=    load_dataset('''hf-internal-testing/fixtures_ade20k'''						,split='''test'''				)
             __SCREAMING_SNAKE_CASE  :Optional[int]													=    Image.open(ds[0]['''file''']				)
             __SCREAMING_SNAKE_CASE  :Union[str, Any]													=    image_processor(images=A_						,return_tensors='''pt'''				).to(A_				)

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

             # verify the logits
             __SCREAMING_SNAKE_CASE  :List[str]													=    torch.Size((1, 1_50, 1_60, 1_60)				)
             self.assertEqual(logits.shape						,A_				)

             __SCREAMING_SNAKE_CASE  :Union[str, Any]													=    version.parse(PIL.__version__				) < version.parse('''9.0.0'''				)

             if is_pillow_less_than_a:
                   __SCREAMING_SNAKE_CASE  :List[str]													=    torch.tensor(
                       [
                           [[-4.9_2_2_5, -2.3_9_5_4, -3.0_5_2_2], [-2.8_8_2_2, -1.0_0_4_6, -1.7_5_6_1], [-2.9_5_4_9, -1.3_2_2_8, -2.1_3_4_7]],
                           [[-5.8_1_6_8, -3.4_1_2_9, -4.0_7_7_8], [-3.8_6_5_1, -2.2_2_1_4, -3.0_2_7_7], [-3.8_3_5_6, -2.4_6_4_3, -3.3_5_3_5]],
                           [[-0.0_0_7_8, 3.9_9_5_2, 4.0_7_5_4], [2.9_8_5_6, 4.6_9_4_4, 5.0_0_3_5], [3.2_4_1_3, 4.7_8_1_3, 4.9_9_6_9]],
                       ]						,device=A_						,)
             else:
                   __SCREAMING_SNAKE_CASE  :List[Any]													=    torch.tensor(
                       [
                           [[-4.8_9_6_0, -2.3_6_8_8, -3.0_3_5_5], [-2.8_4_7_8, -0.9_8_3_6, -1.7_4_1_8], [-2.9_4_4_9, -1.3_3_3_2, -2.1_4_5_6]],
                           [[-5.8_0_8_1, -3.4_1_2_4, -4.1_0_0_6], [-3.8_5_6_1, -2.2_0_8_1, -3.0_3_2_3], [-3.8_3_6_5, -2.4_6_0_1, -3.3_6_6_9]],
                           [[-0.0_3_0_9, 3.9_8_6_8, 4.0_5_4_0], [2.9_6_4_0, 4.6_8_7_7, 4.9_9_7_6], [3.2_0_8_1, 4.7_6_9_0, 4.9_9_4_2]],
                       ]						,device=A_						,)

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




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE  :int													=    BeitForSemanticSegmentation.from_pretrained('''microsoft/beit-base-finetuned-ade-640-640'''				)
             __SCREAMING_SNAKE_CASE  :Optional[int]													=    model.to(A_				)

             __SCREAMING_SNAKE_CASE  :str													=    BeitImageProcessor(do_resize=A_						,size=6_40						,do_center_crop=A_				)

             __SCREAMING_SNAKE_CASE  :Union[str, Any]													=    load_dataset('''hf-internal-testing/fixtures_ade20k'''						,split='''test'''				)
             __SCREAMING_SNAKE_CASE  :int													=    Image.open(ds[0]['''file''']				)
             __SCREAMING_SNAKE_CASE  :Tuple													=    image_processor(images=A_						,return_tensors='''pt'''				).to(A_				)

             # forward pass
             with torch.no_grad():
                   __SCREAMING_SNAKE_CASE  :str													=    model(**A_				)

             __SCREAMING_SNAKE_CASE  :Tuple													=    outputs.logits.detach().cpu()

             __SCREAMING_SNAKE_CASE  :Any													=    image_processor.post_process_semantic_segmentation(outputs=A_						,target_sizes=[(5_00, 3_00)]				)
             __SCREAMING_SNAKE_CASE  :Tuple													=    torch.Size((5_00, 3_00)				)
             self.assertEqual(segmentation[0].shape						,A_				)

             __SCREAMING_SNAKE_CASE  :Optional[Any]													=    image_processor.post_process_semantic_segmentation(outputs=A_				)
             __SCREAMING_SNAKE_CASE  :Tuple													=    torch.Size((1_60, 1_60)				)
             self.assertEqual(segmentation[0].shape						,A_				)



"""simple docstring"""

import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import BatchEncoding, PreTrainedTokenizer
from ...utils import logging


UpperCAmelCase__					=logging.get_logger(__name__)

UpperCAmelCase__					="▁"

UpperCAmelCase__					={
    "vocab_file": "vocab.json",
    "spm_file": "sentencepiece.bpe.model",
    "tokenizer_config_file": "tokenizer_config.json",
}

UpperCAmelCase__					={
    "vocab_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json",
    },
    "spm_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model",
    },
    "tokenizer_config_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json",
    },
}

UpperCAmelCase__					={
    "facebook/m2m100_418M": 1024,
}

# fmt: off
UpperCAmelCase__					={
    "m2m100": ["af", "am", "ar", "ast", "az", "ba", "be", "bg", "bn", "br", "bs", "ca", "ceb", "cs", "cy", "da", "de", "el", "en", "es", "et", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gu", "ha", "he", "hi", "hr", "ht", "hu", "hy", "id", "ig", "ilo", "is", "it", "ja", "jv", "ka", "kk", "km", "kn", "ko", "lb", "lg", "ln", "lo", "lt", "lv", "mg", "mk", "ml", "mn", "mr", "ms", "my", "ne", "nl", "no", "ns", "oc", "or", "pa", "pl", "ps", "pt", "ro", "ru", "sd", "si", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv", "sw", "ta", "th", "tl", "tn", "tr", "uk", "ur", "uz", "vi", "wo", "xh", "yi", "yo", "zh", "zu"],
    "wmt21": ["en", "ha", "is", "ja", "cs", "ru", "zh", "de"]
}





class  lowerCamelCase__    (		_a     ):
      a			:  str           =		VOCAB_FILES_NAMES
      a			:  Dict           =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      a			:  int           =		PRETRAINED_VOCAB_FILES_MAP
      a			:  Union[str, Any]           =		["""input_ids""", """attention_mask"""]

      a			:  List[int]           =		[]
      a			:  List[int]           =		[]
      def __init__(		self						:						Optional[Any]   ,							A_						:						str   ,							A_						:						Dict   ,							A_						:						str=None   ,							A_						:						Dict=None   ,							A_						:						str="<s>"   ,							A_						:						Any="</s>"   ,							A_						:						List[Any]="</s>"   ,							A_						:						List[str]="<pad>"   ,							A_						:						Optional[int]="<unk>"   ,							A_						:						str="m2m100"   ,							A_						:						Optional[Dict[str, Any]] = None   ,							A_						:						Tuple=8   ,							**A_						:						Dict   ,							):






       '''simple docstring'''


       __lowercase						=						{} if sp_model_kwargs is None else sp_model_kwargs

       __lowercase						=						language_codes
       __lowercase						=						FAIRSEQ_LANGUAGE_CODES[language_codes]
       __lowercase						=						{lang_code: F'''__{lang_code}__''' for lang_code in fairseq_language_code}

       __lowercase						=						kwargs.get("""additional_special_tokens"""   ,							[]				)
       kwargs["additional_special_tokens"] += [
           self.get_lang_token(A_				)
           for lang_code in fairseq_language_code
           if self.get_lang_token(A_				) not in kwargs["additional_special_tokens"]
       ]

       super().__init__(
           src_lang=A_   ,							tgt_lang=A_   ,							bos_token=A_   ,							eos_token=A_   ,							sep_token=A_   ,							unk_token=A_   ,							pad_token=A_   ,							language_codes=A_   ,							sp_model_kwargs=self.sp_model_kwargs   ,							num_madeup_words=A_   ,							**A_   ,							)

       __lowercase						=						vocab_file
       __lowercase						=						load_json(A_				)
       __lowercase						=						{v: k for k, v in self.encoder.items()}
       __lowercase						=						spm_file
       __lowercase						=						load_spm(A_   ,							self.sp_model_kwargs				)

       __lowercase						=						len(self.encoder				)

       __lowercase						=						{
           self.get_lang_token(A_				): self.encoder_size + i for i, lang_code in enumerate(A_				)
       }
       __lowercase						=						{lang_code: self.encoder_size + i for i, lang_code in enumerate(A_				)}
       __lowercase						=						{v: k for k, v in self.lang_token_to_id.items()}

       __lowercase						=						src_lang if src_lang is not None else """en"""
       __lowercase						=						tgt_lang
       __lowercase						=						self.get_lang_id(self._src_lang				)
       self.set_src_lang_special_tokens(self._src_lang				)

       __lowercase						=						num_madeup_words
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]				):






       '''simple docstring'''


       return len(self.encoder				) + len(self.lang_token_to_id				)
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Dict				):






       '''simple docstring'''


       return self._src_lang
      @src_lang.setter
      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						str				):






       '''simple docstring'''


       __lowercase						=						new_src_lang
       self.set_src_lang_special_tokens(self._src_lang				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						str				):






       '''simple docstring'''


       return self.sp_model.encode(A_   ,							out_type=A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]   ,							A_						:						str				):






       '''simple docstring'''


       if token in self.lang_token_to_id:
        return self.lang_token_to_id[token]
       return self.encoder.get(A_   ,							self.encoder[self.unk_token]				)
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]   ,							A_						:						int				):






       '''simple docstring'''


       if index in self.id_to_lang_token:
        return self.id_to_lang_token[index]
       return self.decoder.get(A_   ,							self.unk_token				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						Union[str, Any]				):






       '''simple docstring'''


       __lowercase						=						[]
       __lowercase						=						""""""
       for token in tokens:
        # make sure that special tokens are not decoded using sentencepiece model
        if token in self.all_special_tokens:
         out_string += self.sp_model.decode(A_				) + token
         __lowercase						=						[]
        else:
         current_sub_tokens.append(A_				)
       out_string += self.sp_model.decode(A_				)
       return out_string.strip()
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]   ,							A_						:						List[int]   ,							A_						:						Optional[List[int]] = None   ,							A_						:						bool = False				):






       '''simple docstring'''


       if already_has_special_tokens:
        return super().get_special_tokens_mask(
            token_ids_a=A_   ,							token_ids_a=A_   ,							already_has_special_tokens=A_				)

       __lowercase						=						[1] * len(self.prefix_tokens				)
       __lowercase						=						[1] * len(self.suffix_tokens				)
       if token_ids_a is None:
        return prefix_ones + ([0] * len(A_				)) + suffix_ones
       return prefix_ones + ([0] * len(A_				)) + ([0] * len(A_				)) + suffix_ones
      def SCREAMING_SNAKE_CASE_						(		self						:						Any   ,							A_						:						List[int]   ,							A_						:						Optional[List[int]] = None				):






       '''simple docstring'''


       if token_ids_a is None:
        return self.prefix_tokens + token_ids_a + self.suffix_tokens
       # We don't expect to process pairs, but leave the pair logic for API consistency
       return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       __lowercase						=						{self.convert_ids_to_tokens(A_				): i for i in range(self.vocab_size				)}
       vocab.update(self.added_tokens_encoder				)
       return vocab
      def __getstate__(		self						:						List[str]				):






       '''simple docstring'''


       __lowercase						=						self.__dict__.copy()
       __lowercase						=						None
       return state
      def __setstate__(		self						:						Optional[Any]   ,							A_						:						Dict				):






       '''simple docstring'''


       __lowercase						=						d

       # for backward compatibility
       if not hasattr(self   ,							"""sp_model_kwargs"""				):
        __lowercase						=						{}

       __lowercase						=						load_spm(self.spm_file   ,							self.sp_model_kwargs				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						str   ,							A_						:						Optional[str] = None				):






       '''simple docstring'''


       __lowercase						=						Path(A_				)
       if not save_dir.is_dir():
        raise OSError(F'''{save_directory} should be a directory'''				)
       __lowercase						=						save_dir / (
           (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""vocab_file"""]
       )
       __lowercase						=						save_dir / (
           (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""spm_file"""]
       )

       save_json(self.encoder   ,							A_				)

       if os.path.abspath(self.spm_file				) != os.path.abspath(A_				) and os.path.isfile(self.spm_file				):
        copyfile(self.spm_file   ,							A_				)
       elif not os.path.isfile(self.spm_file				):
        with open(A_   ,							"""wb"""				) as fi:
         __lowercase						=						self.sp_model.serialized_model_proto()
         fi.write(A_				)

       return (str(A_				), str(A_				))
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]   ,							A_						:						List[str]   ,							A_						:						str = "en"   ,							A_						:						Optional[List[str]] = None   ,							A_						:						str = "ro"   ,							**A_						:						List[Any]   ,							):






       '''simple docstring'''


       __lowercase						=						src_lang
       __lowercase						=						tgt_lang
       self.set_src_lang_special_tokens(self.src_lang				)
       return super().prepare_seqaseq_batch(A_   ,							A_   ,							**A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]   ,							A_						:						List[Any]   ,							A_						:						Optional[str]   ,							A_						:						Optional[str]   ,							**A_						:						Optional[Any]				):






       '''simple docstring'''


       if src_lang is None or tgt_lang is None:
        raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model"""				)
       __lowercase						=						src_lang
       __lowercase						=						self(A_   ,							add_special_tokens=A_   ,							**A_				)
       __lowercase						=						self.get_lang_id(A_				)
       __lowercase						=						tgt_lang_id
       return inputs
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       self.set_src_lang_special_tokens(self.src_lang				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       self.set_tgt_lang_special_tokens(self.tgt_lang				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]   ,							A_						:						str				):






       '''simple docstring'''


       __lowercase						=						self.get_lang_token(A_				)
       __lowercase						=						self.lang_token_to_id[lang_token]
       __lowercase						=						[self.cur_lang_id]
       __lowercase						=						[self.eos_token_id]
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]   ,							A_						:						str				):






       '''simple docstring'''


       __lowercase						=						self.get_lang_token(A_				)
       __lowercase						=						self.lang_token_to_id[lang_token]
       __lowercase						=						[self.cur_lang_id]
       __lowercase						=						[self.eos_token_id]
      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						str				):






       '''simple docstring'''


       return self.lang_code_to_token[lang]






      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						str				):






       '''simple docstring'''


       __lowercase						=						self.get_lang_token(A_				)
       return self.lang_token_to_id[lang_token]






def lowerCAmelCase_   (					UpperCamelCase__						:			str					, UpperCamelCase__						:			Dict[str, Any]			):






 """simple docstring"""







 __lowercase						=						sentencepiece.SentencePieceProcessor(**UpperCamelCase__			)
 spm.Load(str(UpperCamelCase__			)			)
 return spm


def lowerCAmelCase_   (					UpperCamelCase__						:			str			):






 """simple docstring"""







 with open(UpperCamelCase__					, """r"""			) as f:
  return json.load(UpperCamelCase__			)


def lowerCAmelCase_   (					UpperCamelCase__						:			List[str]					, UpperCamelCase__						:			str			):






 """simple docstring"""







 with open(UpperCamelCase__					, """w"""			) as f:
  json.dump(UpperCamelCase__					, UpperCamelCase__					, indent=2			)
