
"""simple docstring"""

import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

# Register SEW's fairseq modules
from sew_asapp import tasks  # noqa: F401

from transformers import (
    SEWConfig,
    SEWForCTC,
    SEWModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE_            =    logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_            =    {
    '''post_extract_proj''': '''feature_projection''',
    '''encoder.pos_conv.0''': '''encoder.pos_conv_embed.conv''',
    '''self_attn.k_proj''': '''encoder.layers.*.attention.k_proj''',
    '''self_attn.v_proj''': '''encoder.layers.*.attention.v_proj''',
    '''self_attn.q_proj''': '''encoder.layers.*.attention.q_proj''',
    '''self_attn.out_proj''': '''encoder.layers.*.attention.out_proj''',
    '''self_attn_layer_norm''': '''encoder.layers.*.layer_norm''',
    '''fc1''': '''encoder.layers.*.feed_forward.intermediate_dense''',
    '''fc2''': '''encoder.layers.*.feed_forward.output_dense''',
    '''final_layer_norm''': '''encoder.layers.*.final_layer_norm''',
    '''encoder.upsample.0''': '''encoder.upsample.projection''',
    '''encoder.layer_norm''': '''encoder.layer_norm''',
    '''w2v_model.layer_norm''': '''layer_norm''',
    '''w2v_encoder.proj''': '''lm_head''',
    '''mask_emb''': '''masked_spec_embed''',
}




def      lowercase__    (       lowerCAmelCase					:       Tuple			,					lowerCAmelCase					:       Any			,					lowerCAmelCase					:       Union[str, Any]			,					lowerCAmelCase					:       Union[str, Any]			,					lowerCAmelCase					:       Union[str, Any]			)			->   Union[str, Any]:






    """simple docstring"""




    for attribute in key.split('.'			):
        UpperCAmelCase          =       getattr(lowerCAmelCase			,					lowerCAmelCase			)

    if weight_type is not None:
        UpperCAmelCase          =       getattr(lowerCAmelCase			,					lowerCAmelCase			).shape
    else:
        UpperCAmelCase          =       hf_pointer.shape

    assert hf_shape == value.shape, (
        F"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
        F" {value.shape} for {full_name}"
    )

    if weight_type == "weight":
        UpperCAmelCase          =       value
    elif weight_type == "weight_g":
        UpperCAmelCase          =       value
    elif weight_type == "weight_v":
        UpperCAmelCase          =       value
    elif weight_type == "bias":
        UpperCAmelCase          =       value
    else:
        UpperCAmelCase          =       value

    logger.info(F"{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}."			)




def      lowercase__    (       lowerCAmelCase					:       Optional[Any]			,					lowerCAmelCase					:       Optional[int]			,					lowerCAmelCase					:       List[str]			)			->   Optional[Any]:






    """simple docstring"""




    UpperCAmelCase          =       []
    UpperCAmelCase          =       fairseq_model.state_dict()

    UpperCAmelCase          =       hf_model.sew.feature_extractor if is_finetuned else hf_model.feature_extractor

    for name, value in fairseq_dict.items():
        UpperCAmelCase          =       False
        if "conv_layers" in name:
            load_conv_layer(
                lowerCAmelCase			,					lowerCAmelCase			,					lowerCAmelCase			,					lowerCAmelCase			,					hf_model.config.feat_extract_norm == 'group'			,					)
            UpperCAmelCase          =       True
        else:
            for key, mapped_key in MAPPING.items():
                UpperCAmelCase          =       "sew." + mapped_key if (is_finetuned and mapped_key != "lm_head") else mapped_key

                if key in name or key.split('w2v_model.'			)[-1] == name.split('.'			)[0]:
                    UpperCAmelCase          =       True
                    if "*" in mapped_key:
                        UpperCAmelCase          =       name.split(lowerCAmelCase			)[0].split('.'			)[-2]
                        UpperCAmelCase          =       mapped_key.replace('*'			,					lowerCAmelCase			)
                    if "weight_g" in name:
                        UpperCAmelCase          =       "weight_g"
                    elif "weight_v" in name:
                        UpperCAmelCase          =       "weight_v"
                    elif "weight" in name:
                        UpperCAmelCase          =       "weight"
                    elif "bias" in name:
                        UpperCAmelCase          =       "bias"
                    else:
                        UpperCAmelCase          =       None
                    set_recursively(lowerCAmelCase			,					lowerCAmelCase			,					lowerCAmelCase			,					lowerCAmelCase			,					lowerCAmelCase			)
                continue
        if not is_used:
            unused_weights.append(lowerCAmelCase			)

    logger.warning(F"Unused weights: {unused_weights}"			)




def      lowercase__    (       lowerCAmelCase					:       List[Any]			,					lowerCAmelCase					:       List[Any]			,					lowerCAmelCase					:       str			,					lowerCAmelCase					:       List[Any]			,					lowerCAmelCase					:       Tuple			)			->   Tuple:






    """simple docstring"""




    UpperCAmelCase          =       full_name.split('conv_layers.'			)[-1]
    UpperCAmelCase          =       name.split('.'			)
    UpperCAmelCase          =       int(items[0]			)
    UpperCAmelCase          =       int(items[1]			)

    if type_id == 0:
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
                F"{full_name} has size {value.shape}, but"
                F" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."
            )
            UpperCAmelCase          =       value
            logger.info(F"Feat extract conv layer {layer_id} was initialized from {full_name}."			)
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
                F"{full_name} has size {value.shape}, but"
                F" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."
            )
            UpperCAmelCase          =       value
            logger.info(F"Feat extract conv layer {layer_id} was initialized from {full_name}."			)
    elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
                F"{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was"
                " found."
            )
            UpperCAmelCase          =       value
            logger.info(F"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."			)
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
                F"{full_name} has size {value.shape}, but"
                F" {feature_extractor[layer_id].layer_norm.weight.data.shape} was found."
            )
            UpperCAmelCase          =       value
            logger.info(F"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."			)
    else:
        unused_weights.append(lowerCAmelCase			)




def      lowercase__    (       lowerCAmelCase					:       Dict			,					lowerCAmelCase					:       Optional[int]			)			->   Optional[Any]:






    """simple docstring"""




    UpperCAmelCase          =       SEWConfig()
    if is_finetuned:
        UpperCAmelCase          =       model.wav_encoder.wav_model.cfg
    else:
        UpperCAmelCase          =       model.cfg

    UpperCAmelCase          =       fs_config.conv_bias
    UpperCAmelCase          =       eval(fs_config.conv_feature_layers			)
    UpperCAmelCase          =       [x[0] for x in conv_layers]
    UpperCAmelCase          =       [x[1] for x in conv_layers]
    UpperCAmelCase          =       [x[2] for x in conv_layers]
    UpperCAmelCase          =       "gelu"
    UpperCAmelCase          =       "layer" if fs_config.extractor_mode == "layer_norm" else "group"
    UpperCAmelCase          =       0.0
    UpperCAmelCase          =       fs_config.activation_fn.name
    UpperCAmelCase          =       fs_config.encoder_embed_dim
    UpperCAmelCase          =       0.02
    UpperCAmelCase          =       fs_config.encoder_ffn_embed_dim
    UpperCAmelCase          =       1E-5
    UpperCAmelCase          =       fs_config.encoder_layerdrop
    UpperCAmelCase          =       fs_config.encoder_attention_heads
    UpperCAmelCase          =       fs_config.conv_pos_groups
    UpperCAmelCase          =       fs_config.conv_pos
    UpperCAmelCase          =       len(lowerCAmelCase			)
    UpperCAmelCase          =       fs_config.encoder_layers
    UpperCAmelCase          =       fs_config.squeeze_factor

    # take care of any params that are overridden by the Wav2VecCtc model
    if is_finetuned:
        UpperCAmelCase          =       model.cfg
        UpperCAmelCase          =       fs_config.final_dropout
        UpperCAmelCase          =       fs_config.layerdrop
    UpperCAmelCase          =       fs_config.activation_dropout
    UpperCAmelCase          =       fs_config.mask_prob > 0 or fs_config.mask_channel_prob > 0
    UpperCAmelCase          =       fs_config.attention_dropout
    UpperCAmelCase          =       fs_config.dropout_input
    UpperCAmelCase          =       fs_config.dropout
    UpperCAmelCase          =       fs_config.mask_channel_length
    UpperCAmelCase          =       fs_config.mask_channel_prob
    UpperCAmelCase          =       fs_config.mask_length
    UpperCAmelCase          =       fs_config.mask_prob

    UpperCAmelCase          =       "Wav2Vec2FeatureExtractor"
    UpperCAmelCase          =       "Wav2Vec2CTCTokenizer"

    return config




@torch.no_grad()
def      lowercase__    (       lowerCAmelCase					:       int			,					lowerCAmelCase					:       Dict			,					lowerCAmelCase					:       List[Any]=None			,					lowerCAmelCase					:       Union[str, Any]=None			,					lowerCAmelCase					:       List[str]=True			)			->   List[str]:






    """simple docstring"""




    if is_finetuned:
        UpperCAmelCase          =       fairseq.checkpoint_utils.load_model_ensemble_and_task(
            [checkpoint_path]			,					arg_overrides={'data': '/'.join(dict_path.split('/'			)[:-1]			)}			)
    else:
        UpperCAmelCase          =       fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]			)

    if config_path is not None:
        UpperCAmelCase          =       SEWConfig.from_pretrained(lowerCAmelCase			)
    else:
        UpperCAmelCase          =       convert_config(model[0]			,					lowerCAmelCase			)
    UpperCAmelCase          =       model[0].eval()

    UpperCAmelCase          =       True if config.feat_extract_norm == "layer" else False
    UpperCAmelCase          =       WavaVecaFeatureExtractor(
        feature_size=1			,					sampling_rate=16_000			,					padding_value=0			,					do_normalize=lowerCAmelCase			,					return_attention_mask=lowerCAmelCase			,					)

    if is_finetuned:
        if dict_path:
            UpperCAmelCase          =       Dictionary.load(lowerCAmelCase			)

            # important change bos & pad token id since CTC symbol is <pad> and
            # not <s> as in fairseq
            UpperCAmelCase          =       target_dict.pad_index
            UpperCAmelCase          =       target_dict.bos_index
            UpperCAmelCase          =       target_dict.pad_index
            UpperCAmelCase          =       target_dict.bos_index
            UpperCAmelCase          =       target_dict.eos_index
            UpperCAmelCase          =       len(target_dict.symbols			)
            UpperCAmelCase          =       os.path.join(lowerCAmelCase			,					'vocab.json'			)
            if not os.path.isdir(lowerCAmelCase			):
                logger.error('--pytorch_dump_folder_path ({}) should be a directory'.format(lowerCAmelCase			)			)
                return
            os.makedirs(lowerCAmelCase			,					exist_ok=lowerCAmelCase			)
            with open(lowerCAmelCase			,					'w'			,					encoding='utf-8'			) as vocab_handle:
                json.dump(target_dict.indices			,					lowerCAmelCase			)
            UpperCAmelCase          =       WavaVecaCTCTokenizer(
                lowerCAmelCase			,					unk_token=target_dict.unk_word			,					pad_token=target_dict.pad_word			,					bos_token=target_dict.bos_word			,					eos_token=target_dict.eos_word			,					word_delimiter_token='|'			,					do_lower_case=lowerCAmelCase			,					)
            UpperCAmelCase          =       WavaVecaProcessor(feature_extractor=lowerCAmelCase			,					tokenizer=lowerCAmelCase			)
            processor.save_pretrained(lowerCAmelCase			)

        UpperCAmelCase          =       SEWForCTC(lowerCAmelCase			)
    else:
        UpperCAmelCase          =       SEWModel(lowerCAmelCase			)
        feature_extractor.save_pretrained(lowerCAmelCase			)

    recursively_load_weights(lowerCAmelCase			,					lowerCAmelCase			,					lowerCAmelCase			)

    hf_model.save_pretrained(lowerCAmelCase			)


if __name__ == "__main__":
   SCREAMING_SNAKE_CASE_            =    argparse.ArgumentParser()
   parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
   parser.add_argument('''--checkpoint_path''', default=None, type=str, help='''Path to fairseq checkpoint''')
   parser.add_argument('''--dict_path''', default=None, type=str, help='''Path to dict of fine-tuned model''')
   parser.add_argument('''--config_path''', default=None, type=str, help='''Path to hf config.json of model to convert''')
   parser.add_argument(
       '''--is_finetuned''', action='''store_true''', help='''Whether the model to convert is a fine-tuned model or not'''
   )
   SCREAMING_SNAKE_CASE_            =    parser.parse_args()
   convert_sew_checkpoint(
       args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, args.is_finetuned
   )







'''simple docstring'''


import inspect
import unittest

import numpy as np

from transformers import BeitConfig
from transformers.testing_utils import require_flax, require_vision, slow
from transformers.utils import cached_property, is_flax_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor


if is_flax_available():
      import jax

      from transformers import FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling, FlaxBeitModel

if is_vision_available():
      from PIL import Image

      from transformers import BeitImageProcessor







class 							__magic_name__	(       unittest.TestCase ):

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase			,     lowerCamelCase=100			,     lowerCamelCase=13			,     lowerCamelCase=30			,     lowerCamelCase=2			,     lowerCamelCase=3			,     lowerCamelCase=True			,     lowerCamelCase=True			,     lowerCamelCase=32			,     lowerCamelCase=5			,     lowerCamelCase=4			,     lowerCamelCase=37			,     lowerCamelCase="gelu"			,     lowerCamelCase=0.1			,     lowerCamelCase=0.1			,     lowerCamelCase=10			,     lowerCamelCase=0.02			,     lowerCamelCase=3			,     ):


          '''simple docstring'''





          __A	:       Tuple      =							parent
          __A	:       Dict      =							vocab_size
          __A	:       Union[str, Any]      =							batch_size
          __A	:       str      =							image_size
          __A	:       Optional[Any]      =							patch_size
          __A	:       Optional[Any]      =							num_channels
          __A	:       Optional[Any]      =							is_training
          __A	:       Any      =							use_labels
          __A	:       List[str]      =							hidden_size
          __A	:       Union[str, Any]      =							num_hidden_layers
          __A	:       Optional[Any]      =							num_attention_heads
          __A	:       Optional[int]      =							intermediate_size
          __A	:       Optional[int]      =							hidden_act
          __A	:       Dict      =							hidden_dropout_prob
          __A	:       str      =							attention_probs_dropout_prob
          __A	:       Optional[int]      =							type_sequence_label_size
          __A	:       Optional[int]      =							initializer_range

          # 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	:       Union[str, Any]      =							num_patches + 1


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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

          __A	:       str      =							None
          if self.use_labels:
               __A	:       List[str]      =							ids_tensor([self.batch_size]			,     self.type_sequence_label_size		)

          __A	:       Dict      =							BeitConfig(
              vocab_size=self.vocab_size			,     image_size=self.image_size			,     patch_size=self.patch_size			,     num_channels=self.num_channels			,     hidden_size=self.hidden_size			,     num_hidden_layers=self.num_hidden_layers			,     num_attention_heads=self.num_attention_heads			,     intermediate_size=self.intermediate_size			,     hidden_act=self.hidden_act			,     hidden_dropout_prob=self.hidden_dropout_prob			,     attention_probs_dropout_prob=self.attention_probs_dropout_prob			,     is_decoder=lowerCamelCase			,     initializer_range=self.initializer_range			,     )

          return config, pixel_values, labels


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Optional[int]      =							FlaxBeitModel(config=lowerCamelCase		)
          __A	:       Optional[Any]      =							model(lowerCamelCase		)
          self.parent.assertEqual(result.last_hidden_state.shape			,     (self.batch_size, self.seq_length, self.hidden_size)		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Dict      =							FlaxBeitForMaskedImageModeling(config=lowerCamelCase		)
          __A	:       Union[str, Any]      =							model(lowerCamelCase		)
          self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.seq_length - 1, self.vocab_size)		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Any      =							self.type_sequence_label_size
          __A	:       List[str]      =							FlaxBeitForImageClassification(config=lowerCamelCase		)
          __A	:       Tuple      =							model(lowerCamelCase		)
          self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.type_sequence_label_size)		)

          # test greyscale images
          __A	:       Tuple      =							1
          __A	:       Optional[Any]      =							FlaxBeitForImageClassification(lowerCamelCase		)

          __A	:       Union[str, Any]      =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
          __A	:       str      =							model(lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Tuple      =							self.prepare_config_and_inputs()
          (
              (
              __A
          )  ,(
              __A
          )  ,(
              __A
          )  ,
          )	:       Dict      =							config_and_inputs
          __A	:       Dict      =							{"pixel_values": pixel_values}
          return config, inputs_dict







@require_flax
class 							__magic_name__	(       lowerCAmelCase    ,		unittest.TestCase ):

     """simple docstring"""






     lowerCamelCase__       						=				(
         (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
     )


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[int]      =							FlaxBeitModelTester(self		)
          __A	:       List[Any]      =							ConfigTester(self			,     config_class=lowerCamelCase			,     has_text_modality=lowerCamelCase			,     hidden_size=37		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          self.config_tester.run_common_tests()


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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

          for model_class in self.all_model_classes:
               __A	:       int      =							model_class(lowerCamelCase		)
               __A	:       Dict      =							inspect.signature(model.__call__		)
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               __A	:       List[str]      =							[*signature.parameters.keys()]

               __A	:       Optional[Any]      =							["pixel_values"]
               self.assertListEqual(arg_names[:1]			,     lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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

          for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__		):
                    __A	:       Optional[int]      =							self._prepare_for_class(lowerCamelCase			,     lowerCamelCase		)
                    __A	:       Optional[int]      =							model_class(lowerCamelCase		)

                    @jax.jit
                    def model_jitted(lowerCamelCase			,     **lowerCamelCase		):
                         return model(pixel_values=lowerCamelCase			,     **lowerCamelCase		)

                    with self.subTest("JIT Enabled"		):
                         __A	:       Tuple      =							model_jitted(**lowerCamelCase		).to_tuple()

                    with self.subTest("JIT Disabled"		):
                         with jax.disable_jit():
                              __A	:       Dict      =							model_jitted(**lowerCamelCase		).to_tuple()

                    self.assertEqual(len(lowerCamelCase		)			,     len(lowerCamelCase		)		)
                    for jitted_output, output in zip(lowerCamelCase			,     lowerCamelCase		):
                         self.assertEqual(jitted_output.shape			,     output.shape		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Dict      =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Optional[Any]      =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       List[str]      =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*lowerCamelCase		)


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          for model_class_name in self.all_model_classes:
               __A	:       Any      =							model_class_name.from_pretrained("microsoft/beit-base-patch16-224"		)
               __A	:       Optional[int]      =							model(np.ones((1, 3, 224, 224)		)		)
               self.assertIsNotNone(lowerCamelCase		)


def 		_lowercase ():




     '''simple docstring'''





     __A	:       Dict      =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"							)
     return image







@require_vision
@require_flax
class 							__magic_name__	(       unittest.TestCase ):

     """simple docstring"""






     @cached_property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Union[str, Any]      =							FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"		)

          __A	:       List[Any]      =							self.default_image_processor
          __A	:       Dict      =							prepare_img()
          __A	:       str      =							image_processor(images=lowerCamelCase			,     return_tensors="np"		).pixel_values

          # prepare bool_masked_pos
          __A	:       List[Any]      =							np.ones((1, 196)			,     dtype=lowerCamelCase		)

          # forward pass
          __A	:       Union[str, Any]      =							model(pixel_values=lowerCamelCase			,     bool_masked_pos=lowerCamelCase		)
          __A	:       Dict      =							outputs.logits

          # verify the logits
          __A	:       int      =							(1, 196, 8192)
          self.assertEqual(logits.shape			,     lowerCamelCase		)

          __A	:       Optional[int]      =							np.array(
              [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]		)

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


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Union[str, Any]      =							FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"		)

          __A	:       Any      =							self.default_image_processor
          __A	:       List[Any]      =							prepare_img()
          __A	:       List[Any]      =							image_processor(images=lowerCamelCase			,     return_tensors="np"		)

          # forward pass
          __A	:       int      =							model(**lowerCamelCase		)
          __A	:       Any      =							outputs.logits

          # verify the logits
          __A	:       str      =							(1, 1000)
          self.assertEqual(logits.shape			,     lowerCamelCase		)

          __A	:       int      =							np.array([-1.2385, -1.0987, -1.0108]		)

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

          __A	:       Optional[int]      =							281
          self.assertEqual(logits.argmax(-1		).item()			,     lowerCamelCase		)


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Any      =							FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"		)

          __A	:       Optional[int]      =							self.default_image_processor
          __A	:       Tuple      =							prepare_img()
          __A	:       List[str]      =							image_processor(images=lowerCamelCase			,     return_tensors="np"		)

          # forward pass
          __A	:       List[Any]      =							model(**lowerCamelCase		)
          __A	:       Dict      =							outputs.logits

          # verify the logits
          __A	:       Union[str, Any]      =							(1, 2_1841)
          self.assertEqual(logits.shape			,     lowerCamelCase		)

          __A	:       Tuple      =							np.array([1.6881, -0.2787, 0.5901]		)

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

          __A	:       Optional[int]      =							2396
          self.assertEqual(logits.argmax(-1		).item()			,     lowerCamelCase		)

