






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







    """simple docstring"""







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





           '''simple docstring'''

           __magic_name__      				=    parent
           __magic_name__      				=    1_0_0
           __magic_name__      				=    batch_size
           __magic_name__      				=    image_size
           __magic_name__      				=    patch_size
           __magic_name__      				=    num_channels
           __magic_name__      				=    is_training
           __magic_name__      				=    use_labels
           __magic_name__      				=    hidden_size
           __magic_name__      				=    num_hidden_layers
           __magic_name__      				=    num_attention_heads
           __magic_name__      				=    intermediate_size
           __magic_name__      				=    hidden_act
           __magic_name__      				=    hidden_dropout_prob
           __magic_name__      				=    attention_probs_dropout_prob
           __magic_name__      				=    type_sequence_label_size
           __magic_name__      				=    initializer_range
           __magic_name__      				=    scope
           __magic_name__      				=    out_indices
           __magic_name__      				=    num_labels

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







    def 				a__			(		self :	Dict ):





           '''simple docstring'''

           __magic_name__      				=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

           __magic_name__      				=    self.get_config()

           return config, pixel_values, labels, pixel_labels







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







    def 				a__			(		self :	str				,    UpperCamelCase_ :	Dict				,    UpperCamelCase_ :	Any				,    UpperCamelCase_ :	Dict				,    UpperCamelCase_ :	Optional[Any] ):





           '''simple docstring'''

           __magic_name__      				=    BeitModel(config=UpperCamelCase_ )
           model.to(UpperCamelCase_ )
           model.eval()
           __magic_name__      				=    model(UpperCamelCase_ )
           self.parent.assertEqual(result.last_hidden_state.shape				,    (self.batch_size, self.seq_length, self.hidden_size) )







    def 				a__			(		self :	Optional[Any]				,    UpperCamelCase_ :	Tuple				,    UpperCamelCase_ :	List[str]				,    UpperCamelCase_ :	Tuple				,    UpperCamelCase_ :	Any ):





           '''simple docstring'''

           __magic_name__      				=    BeitForMaskedImageModeling(config=UpperCamelCase_ )
           model.to(UpperCamelCase_ )
           model.eval()
           __magic_name__      				=    model(UpperCamelCase_ )
           self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.seq_length - 1, self.vocab_size) )







    def 				a__			(		self :	Any				,    UpperCamelCase_ :	Optional[int]				,    UpperCamelCase_ :	List[Any]				,    UpperCamelCase_ :	List[Any]				,    UpperCamelCase_ :	List[str] ):





           '''simple docstring'''

           __magic_name__      				=    self.type_sequence_label_size
           __magic_name__      				=    BeitForImageClassification(UpperCamelCase_ )
           model.to(UpperCamelCase_ )
           model.eval()
           __magic_name__      				=    model(UpperCamelCase_				,    labels=UpperCamelCase_ )
           self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.type_sequence_label_size) )

           # test greyscale images
           __magic_name__      				=    1
           __magic_name__      				=    BeitForImageClassification(UpperCamelCase_ )
           model.to(UpperCamelCase_ )
           model.eval()

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







    def 				a__			(		self :	List[str]				,    UpperCamelCase_ :	Dict				,    UpperCamelCase_ :	int				,    UpperCamelCase_ :	Dict				,    UpperCamelCase_ :	Dict ):





           '''simple docstring'''

           __magic_name__      				=    self.num_labels
           __magic_name__      				=    BeitForSemanticSegmentation(UpperCamelCase_ )
           model.to(UpperCamelCase_ )
           model.eval()
           __magic_name__      				=    model(UpperCamelCase_ )
           self.parent.assertEqual(
               result.logits.shape				,    (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2) )
           __magic_name__      				=    model(UpperCamelCase_				,    labels=UpperCamelCase_ )
           self.parent.assertEqual(
               result.logits.shape				,    (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2) )







    def 				a__			(		self :	Tuple ):





           '''simple docstring'''

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



@require_torch
class        UpperCamelCase__ (				a_							,	a_							,	unittest.TestCase):







    """simple docstring"""
    __UpperCAmelCase           = (
        (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
        if is_torch_available()
        else ()
    )
    __UpperCAmelCase           = (
        {
            """feature-extraction""": BeitModel,
            """image-classification""": BeitForImageClassification,
            """image-segmentation""": BeitForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    __UpperCAmelCase           = False
    __UpperCAmelCase           = False
    __UpperCAmelCase           = False







    def 				a__			(		self :	Any ):





           '''simple docstring'''

           __magic_name__      				=    BeitModelTester(self )
           __magic_name__      				=    ConfigTester(self				,    config_class=UpperCamelCase_				,    has_text_modality=UpperCamelCase_				,    hidden_size=3_7 )







    def 				a__			(		self :	int ):





           '''simple docstring'''

           self.config_tester.run_common_tests()







    @unittest.skip(reason='BEiT does not use inputs_embeds' )
    def 				a__			(		self :	List[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 				a__			(		self :	Dict ):





           '''simple docstring'''

           pass







    def 				a__			(		self :	Dict ):





           '''simple docstring'''

           __magic_name__				,  __magic_name__      				=    self.model_tester.prepare_config_and_inputs_for_common()

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







    def 				a__			(		self :	Tuple ):





           '''simple docstring'''

           __magic_name__				,  __magic_name__      				=    self.model_tester.prepare_config_and_inputs_for_common()

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

                  __magic_name__      				=    ['pixel_values']
                  self.assertListEqual(arg_names[:1]				,    UpperCamelCase_ )







    def 				a__			(		self :	Optional[int] ):





           '''simple docstring'''

           __magic_name__      				=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase_ )







    def 				a__			(		self :	Union[str, Any] ):





           '''simple docstring'''

           __magic_name__      				=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_ )







    def 				a__			(		self :	Dict ):





           '''simple docstring'''

           __magic_name__      				=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCamelCase_ )







    def 				a__			(		self :	Dict ):





           '''simple docstring'''

           __magic_name__      				=    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*UpperCamelCase_ )







    def 				a__			(		self :	Union[str, Any] ):





           '''simple docstring'''

           if not self.model_tester.is_training:
                  return

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

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

                  __magic_name__      				=    model_class(UpperCamelCase_ )
                  model.to(UpperCamelCase_ )
                  model.train()
                  __magic_name__      				=    self._prepare_for_class(UpperCamelCase_				,    UpperCamelCase_				,    return_labels=UpperCamelCase_ )
                  __magic_name__      				=    model(**UpperCamelCase_ ).loss
                  loss.backward()







    def 				a__			(		self :	Union[str, Any] ):





           '''simple docstring'''

           __magic_name__				,  __magic_name__      				=    self.model_tester.prepare_config_and_inputs_for_common()
           if not self.model_tester.is_training:
                  return

           __magic_name__      				=    False
           __magic_name__      				=    True

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

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







    def 				a__			(		self :	Optional[int] ):





           '''simple docstring'''

           __magic_name__				,  __magic_name__      				=    self.model_tester.prepare_config_and_inputs_for_common()

           __magic_name__      				=    _config_zero_init(UpperCamelCase_ )
           for model_class in self.all_model_classes:
                  __magic_name__      				=    model_class(config=UpperCamelCase_ )
                  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 :	List[str] ):





           '''simple docstring'''

           for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  __magic_name__      				=    BeitModel.from_pretrained(UpperCamelCase_ )
                  self.assertIsNotNone(UpperCamelCase_ )


def 							A						(				)    ->		Any:

       """simple docstring"""





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






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







    """simple docstring"""







    @cached_property
    def 				a__			(		self :	int ):





           '''simple docstring'''

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







    @slow
    def 				a__			(		self :	str ):





           '''simple docstring'''

           __magic_name__      				=    BeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k' ).to(UpperCamelCase_ )

           __magic_name__      				=    self.default_image_processor
           __magic_name__      				=    prepare_img()
           __magic_name__      				=    image_processor(images=UpperCamelCase_				,    return_tensors='pt' ).pixel_values.to(UpperCamelCase_ )

           # prepare bool_masked_pos
           __magic_name__      				=    torch.ones((1, 1_9_6)				,    dtype=torch.bool ).to(UpperCamelCase_ )

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

           # verify the logits
           __magic_name__      				=    torch.Size((1, 1_9_6, 8_1_9_2) )
           self.assertEqual(logits.shape				,    UpperCamelCase_ )

           __magic_name__      				=    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(UpperCamelCase_ )

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







    @slow
    def 				a__			(		self :	Any ):





           '''simple docstring'''

           __magic_name__      				=    BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224' ).to(UpperCamelCase_ )

           __magic_name__      				=    self.default_image_processor
           __magic_name__      				=    prepare_img()
           __magic_name__      				=    image_processor(images=UpperCamelCase_				,    return_tensors='pt' ).to(UpperCamelCase_ )

           # forward pass
           with torch.no_grad():
                  __magic_name__      				=    model(**UpperCamelCase_ )
           __magic_name__      				=    outputs.logits

           # verify the logits
           __magic_name__      				=    torch.Size((1, 1_0_0_0) )
           self.assertEqual(logits.shape				,    UpperCamelCase_ )

           __magic_name__      				=    torch.tensor([-1.2_385, -1.0_987, -1.0_108] ).to(UpperCamelCase_ )

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

           __magic_name__      				=    2_8_1
           self.assertEqual(logits.argmax(-1 ).item()				,    UpperCamelCase_ )







    @slow
    def 				a__			(		self :	int ):





           '''simple docstring'''

           __magic_name__      				=    BeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k' ).to(
               UpperCamelCase_ )

           __magic_name__      				=    self.default_image_processor
           __magic_name__      				=    prepare_img()
           __magic_name__      				=    image_processor(images=UpperCamelCase_				,    return_tensors='pt' ).to(UpperCamelCase_ )

           # forward pass
           with torch.no_grad():
                  __magic_name__      				=    model(**UpperCamelCase_ )
           __magic_name__      				=    outputs.logits

           # verify the logits
           __magic_name__      				=    torch.Size((1, 2_1_8_4_1) )
           self.assertEqual(logits.shape				,    UpperCamelCase_ )

           __magic_name__      				=    torch.tensor([1.6_881, -0.2_787, 0.5_901] ).to(UpperCamelCase_ )

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

           __magic_name__      				=    2_3_9_6
           self.assertEqual(logits.argmax(-1 ).item()				,    UpperCamelCase_ )







    @slow
    def 				a__			(		self :	Union[str, Any] ):





           '''simple docstring'''

           __magic_name__      				=    BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640' )
           __magic_name__      				=    model.to(UpperCamelCase_ )

           __magic_name__      				=    BeitImageProcessor(do_resize=UpperCamelCase_				,    size=6_4_0				,    do_center_crop=UpperCamelCase_ )

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

           # forward pass
           with torch.no_grad():
                  __magic_name__      				=    model(**UpperCamelCase_ )
           __magic_name__      				=    outputs.logits

           # verify the logits
           __magic_name__      				=    torch.Size((1, 1_5_0, 1_6_0, 1_6_0) )
           self.assertEqual(logits.shape				,    UpperCamelCase_ )

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

           if is_pillow_less_than_a:
                  __magic_name__      				=    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=UpperCamelCase_				,    )
           else:
                  __magic_name__      				=    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=UpperCamelCase_				,    )

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







    @slow
    def 				a__			(		self :	int ):





           '''simple docstring'''

           __magic_name__      				=    BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640' )
           __magic_name__      				=    model.to(UpperCamelCase_ )

           __magic_name__      				=    BeitImageProcessor(do_resize=UpperCamelCase_				,    size=6_4_0				,    do_center_crop=UpperCamelCase_ )

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

           # forward pass
           with torch.no_grad():
                  __magic_name__      				=    model(**UpperCamelCase_ )

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

           __magic_name__      				=    image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase_				,    target_sizes=[(5_0_0, 3_0_0)] )
           __magic_name__      				=    torch.Size((5_0_0, 3_0_0) )
           self.assertEqual(segmentation[0].shape				,    UpperCamelCase_ )

           __magic_name__      				=    image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase_ )
           __magic_name__      				=    torch.Size((1_6_0, 1_6_0) )
           self.assertEqual(segmentation[0].shape				,    UpperCamelCase_ )







"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


snake_case : int     		=		{}

try:
 if not is_sentencepiece_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 snake_case : str     		=		["""NllbTokenizer"""]

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 snake_case : Dict     		=		["""NllbTokenizerFast"""]


if TYPE_CHECKING:
 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_nllb import NllbTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_nllb_fast import NllbTokenizerFast

else:
 import sys

 snake_case : Dict     		=		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)