



'''simple docstring'''





import argparse
import collections
import os
import re
import tempfile

import pandas as pd
from datasets import Dataset
from huggingface_hub import hf_hub_download, upload_folder

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/update_metadata.py
SCREAMING_SNAKE_CASE__      					=     "src/transformers"


# This is to make sure the transformers module imported is the one in the repo.
SCREAMING_SNAKE_CASE__      					=     direct_transformers_import(TRANSFORMERS_PATH)


# Regexes that match TF/Flax/PT model names.
SCREAMING_SNAKE_CASE__      					=     re.compile(R"TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
SCREAMING_SNAKE_CASE__      					=     re.compile(R"Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
SCREAMING_SNAKE_CASE__      					=     re.compile(R"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")


# Fill this with tuples (pipeline_tag, model_mapping, auto_model)
SCREAMING_SNAKE_CASE__      					=     [
    ("pretraining", "MODEL_FOR_PRETRAINING_MAPPING_NAMES", "AutoModelForPreTraining"),
    ("feature-extraction", "MODEL_MAPPING_NAMES", "AutoModel"),
    ("audio-classification", "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES", "AutoModelForAudioClassification"),
    ("text-generation", "MODEL_FOR_CAUSAL_LM_MAPPING_NAMES", "AutoModelForCausalLM"),
    ("automatic-speech-recognition", "MODEL_FOR_CTC_MAPPING_NAMES", "AutoModelForCTC"),
    ("image-classification", "MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES", "AutoModelForImageClassification"),
    ("image-segmentation", "MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES", "AutoModelForImageSegmentation"),
    ("fill-mask", "MODEL_FOR_MASKED_LM_MAPPING_NAMES", "AutoModelForMaskedLM"),
    ("object-detection", "MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES", "AutoModelForObjectDetection"),
    (
        "zero-shot-object-detection",
        "MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES",
        "AutoModelForZeroShotObjectDetection",
    ),
    ("question-answering", "MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES", "AutoModelForQuestionAnswering"),
    ("text2text-generation", "MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES", "AutoModelForSeq2SeqLM"),
    ("text-classification", "MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES", "AutoModelForSequenceClassification"),
    ("automatic-speech-recognition", "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES", "AutoModelForSpeechSeq2Seq"),
    (
        "table-question-answering",
        "MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES",
        "AutoModelForTableQuestionAnswering",
    ),
    ("token-classification", "MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES", "AutoModelForTokenClassification"),
    ("multiple-choice", "MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES", "AutoModelForMultipleChoice"),
    (
        "next-sentence-prediction",
        "MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES",
        "AutoModelForNextSentencePrediction",
    ),
    (
        "audio-frame-classification",
        "MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES",
        "AutoModelForAudioFrameClassification",
    ),
    ("audio-xvector", "MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES", "AutoModelForAudioXVector"),
    (
        "document-question-answering",
        "MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES",
        "AutoModelForDocumentQuestionAnswering",
    ),
    (
        "visual-question-answering",
        "MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES",
        "AutoModelForVisualQuestionAnswering",
    ),
    ("image-to-text", "MODEL_FOR_FOR_VISION_2_SEQ_MAPPING_NAMES", "AutoModelForVision2Seq"),
    (
        "zero-shot-image-classification",
        "MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES",
        "AutoModelForZeroShotImageClassification",
    ),
    ("depth-estimation", "MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES", "AutoModelForDepthEstimation"),
    ("video-classification", "MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES", "AutoModelForVideoClassification"),
    ("mask-generation", "MODEL_FOR_MASK_GENERATION_MAPPING_NAMES", "AutoModelForMaskGeneration"),
]




def        lowerCamelCase  (				_snake_case  :						Dict       ):





    '''simple docstring'''


    lowercase__		       =						re.finditer(".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)"     ,_snake_case       )
    return [m.group(0       ) for m in matches]




def        lowerCamelCase  (				):





    '''simple docstring'''


    lowercase__		       =						transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
    lowercase__		       =						{
        config.replace("Config"     ,""       ): model_type for model_type, config in config_maping_names.items()
    }

    # Dictionaries flagging if each model prefix has a backend in PT/TF/Flax.
    lowercase__		       =						collections.defaultdict(_snake_case       )
    lowercase__		       =						collections.defaultdict(_snake_case       )
    lowercase__		       =						collections.defaultdict(_snake_case       )

    # Let's lookup through all transformers object (once) and find if models are supported by a given backend.
    for attr_name in dir(_snake_case       ):
        lowercase__		       =						None
        if _re_tf_models.match(_snake_case       ) is not None:
            lowercase__		       =						tf_models
            lowercase__		       =						_re_tf_models.match(_snake_case       ).groups()[0]
        elif _re_flax_models.match(_snake_case       ) is not None:
            lowercase__		       =						flax_models
            lowercase__		       =						_re_flax_models.match(_snake_case       ).groups()[0]
        elif _re_pt_models.match(_snake_case       ) is not None:
            lowercase__		       =						pt_models
            lowercase__		       =						_re_pt_models.match(_snake_case       ).groups()[0]

        if lookup_dict is not None:
            while len(_snake_case       ) > 0:
                if attr_name in model_prefix_to_model_type:
                    lowercase__		       =						True
                    break
                # Try again after removing the last word in the name
                lowercase__		       =						"".join(camel_case_split(_snake_case       )[:-1]       )

    lowercase__		       =						set(list(pt_models.keys()       ) + list(tf_models.keys()       ) + list(flax_models.keys()       )       )
    lowercase__		       =						list(_snake_case       )
    all_models.sort()

    lowercase__		       =						{"model_type": all_models}
    lowercase__		       =						[pt_models[t] for t in all_models]
    lowercase__		       =						[tf_models[t] for t in all_models]
    lowercase__		       =						[flax_models[t] for t in all_models]

    # Now let's use the auto-mapping names to make sure
    lowercase__		       =						{}
    for t in all_models:
        if t in transformers_module.models.auto.processing_auto.PROCESSOR_MAPPING_NAMES:
            lowercase__		       =						"AutoProcessor"
        elif t in transformers_module.models.auto.tokenization_auto.TOKENIZER_MAPPING_NAMES:
            lowercase__		       =						"AutoTokenizer"
        elif t in transformers_module.models.auto.feature_extraction_auto.FEATURE_EXTRACTOR_MAPPING_NAMES:
            lowercase__		       =						"AutoFeatureExtractor"
        else:
            # Default to AutoTokenizer if a model has nothing, for backward compatibility.
            lowercase__		       =						"AutoTokenizer"

    lowercase__		       =						[processors[t] for t in all_models]

    return pd.DataFrame(_snake_case       )




def        lowerCamelCase  (				_snake_case  :						List[str]       ):





    '''simple docstring'''


    lowercase__		       =						[
        transformers_module.models.auto.modeling_auto,
        transformers_module.models.auto.modeling_tf_auto,
        transformers_module.models.auto.modeling_flax_auto,
    ]
    for pipeline_tag, model_mapping, auto_class in PIPELINE_TAGS_AND_AUTO_MODELS:
        lowercase__		       =						[model_mapping, f'''TF_{model_mapping}''', f'''FLAX_{model_mapping}''']
        lowercase__		       =						[auto_class, f'''TF_{auto_class}''', f'''Flax_{auto_class}''']
        # Loop through all three frameworks
        for module, cls, mapping in zip(_snake_case     ,_snake_case     ,_snake_case       ):
            # The type of pipeline may not exist in this framework
            if not hasattr(_snake_case     ,_snake_case       ):
                continue
            # First extract all model_names
            lowercase__		       =						[]
            for name in getattr(_snake_case     ,_snake_case       ).values():
                if isinstance(_snake_case     ,_snake_case       ):
                    model_names.append(_snake_case       )
                else:
                    model_names.extend(list(_snake_case       )       )

            # Add pipeline tag and auto model class for those models
            table.update({model_name: (pipeline_tag, cls) for model_name in model_names}       )

    return table




def        lowerCamelCase  (				_snake_case  :						str     ,_snake_case  :						str       ):





    '''simple docstring'''


    lowercase__		       =						get_frameworks_table()
    lowercase__		       =						Dataset.from_pandas(_snake_case       )

    lowercase__		       =						hf_hub_download(
        "huggingface/transformers-metadata"     ,"pipeline_tags.json"     ,repo_type="dataset"     ,token=_snake_case       )
    lowercase__		       =						Dataset.from_json(_snake_case       )
    lowercase__		       =						{
        tags_dataset[i]["model_class"]: (tags_dataset[i]["pipeline_tag"], tags_dataset[i]["auto_class"])
        for i in range(len(_snake_case       )       )
    }
    lowercase__		       =						update_pipeline_and_auto_class_table(_snake_case       )

    # Sort the model classes to avoid some nondeterministic updates to create false update commits.
    lowercase__		       =						sorted(table.keys()       )
    lowercase__		       =						pd.DataFrame(
        {
            "model_class": model_classes,
            "pipeline_tag": [table[m][0] for m in model_classes],
            "auto_class": [table[m][1] for m in model_classes],
        }       )
    lowercase__		       =						Dataset.from_pandas(_snake_case       )

    with tempfile.TemporaryDirectory() as tmp_dir:
        frameworks_dataset.to_json(os.path.join(_snake_case     ,"frameworks.json"       )       )
        tags_dataset.to_json(os.path.join(_snake_case     ,"pipeline_tags.json"       )       )

        if commit_sha is not None:
            lowercase__		       =						(
                f'''Update with commit {commit_sha}\n\nSee: '''
                f'''https://github.com/huggingface/transformers/commit/{commit_sha}'''
            )
        else:
            lowercase__		       =						"Update"

        upload_folder(
            repo_id="huggingface/transformers-metadata"     ,folder_path=_snake_case     ,repo_type="dataset"     ,token=_snake_case     ,commit_message=_snake_case     ,)




def        lowerCamelCase  (				):





    '''simple docstring'''


    lowercase__		       =						{tag: cls for tag, _, cls in PIPELINE_TAGS_AND_AUTO_MODELS}
    lowercase__		       =						transformers_module.pipelines.SUPPORTED_TASKS
    lowercase__		       =						[]
    for key in pipeline_tasks:
        if key not in in_table:
            lowercase__		       =						pipeline_tasks[key]["pt"]
            if isinstance(_snake_case     ,(list, tuple)       ):
                lowercase__		       =						model[0]
            lowercase__		       =						model.__name__
            if model not in in_table.values():
                missing.append(_snake_case       )

    if len(_snake_case       ) > 0:
        lowercase__		       =						", ".join(_snake_case       )
        raise ValueError(
            "The following pipeline tags are not present in the `PIPELINE_TAGS_AND_AUTO_MODELS` constant inside "
            f'''`utils/update_metadata.py`: {msg}. Please add them!'''       )


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE__      					=     argparse.ArgumentParser()
 parser.add_argument("--token", type=str, help="The token to use to push to the transformers-metadata dataset.")
 parser.add_argument("--commit_sha", type=str, help="The sha of the commit going with this update.")
 parser.add_argument("--check-only", action="store_true", help="Activate to just check all pipelines are present.")
 SCREAMING_SNAKE_CASE__      					=     parser.parse_args()

 if args.check_only:
  check_pipeline_tags()
 else:
  update_metadata(args.token, args.commit_sha)









'''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 				snake_case (unittest.TestCase   ):



       def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_=100     ,UpperCAmelCase_=13     ,UpperCAmelCase_=30     ,UpperCAmelCase_=2     ,UpperCAmelCase_=3     ,UpperCAmelCase_=True     ,UpperCAmelCase_=True     ,UpperCAmelCase_=32     ,UpperCAmelCase_=5     ,UpperCAmelCase_=4     ,UpperCAmelCase_=37     ,UpperCAmelCase_="gelu"     ,UpperCAmelCase_=0.1     ,UpperCAmelCase_=0.1     ,UpperCAmelCase_=10     ,UpperCAmelCase_=0.02     ,UpperCAmelCase_=3     ,)  ->							List[Any]:
           lowercase__		       =						parent
           lowercase__		       =						vocab_size
           lowercase__		       =						batch_size
           lowercase__		       =						image_size
           lowercase__		       =						patch_size
           lowercase__		       =						num_channels
           lowercase__		       =						is_training
           lowercase__		       =						use_labels
           lowercase__		       =						hidden_size
           lowercase__		       =						num_hidden_layers
           lowercase__		       =						num_attention_heads
           lowercase__		       =						intermediate_size
           lowercase__		       =						hidden_act
           lowercase__		       =						hidden_dropout_prob
           lowercase__		       =						attention_probs_dropout_prob
           lowercase__		       =						type_sequence_label_size
           lowercase__		       =						initializer_range

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



       def 							_a      (				self      )  ->							List[str]:
           lowercase__		       =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

           lowercase__		       =						None
           if self.use_labels:
               lowercase__		       =						ids_tensor([self.batch_size]     ,self.type_sequence_label_size      )

           lowercase__		       =						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     ,)

           return config, pixel_values, labels



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Optional[int]:
           lowercase__		       =						FlaxBeitModel(config=UpperCAmelCase_      )
           lowercase__		       =						model(UpperCAmelCase_      )
           self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.hidden_size)      )



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Tuple:
           lowercase__		       =						FlaxBeitForMaskedImageModeling(config=UpperCAmelCase_      )
           lowercase__		       =						model(UpperCAmelCase_      )
           self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.seq_length - 1, self.vocab_size)      )



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							List[str]:
           lowercase__		       =						self.type_sequence_label_size
           lowercase__		       =						FlaxBeitForImageClassification(config=UpperCAmelCase_      )
           lowercase__		       =						model(UpperCAmelCase_      )
           self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.type_sequence_label_size)      )

           # test greyscale images
           lowercase__		       =						1
           lowercase__		       =						FlaxBeitForImageClassification(UpperCAmelCase_      )

           lowercase__		       =						floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
           lowercase__		       =						model(UpperCAmelCase_      )







       def 							_a      (				self      )  ->							int:
           lowercase__		       =						self.prepare_config_and_inputs()
           (
               (
               lowercase__
           )       ,       (
               lowercase__
           )       ,       (
               lowercase__
           )       ,       
           )		       =						config_and_inputs
           lowercase__		       =						{"pixel_values": pixel_values}
           return config, inputs_dict





@require_flax
class 				snake_case (UpperCamelCase		,      unittest.TestCase   ):
       lowerCAmelCase__  :Optional[int]						   =					(
           (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
       )



       def 							_a      (				self      )  ->							None:
           lowercase__		       =						FlaxBeitModelTester(self      )
           lowercase__		       =						ConfigTester(self     ,config_class=UpperCAmelCase_     ,has_text_modality=UpperCAmelCase_     ,hidden_size=37      )



       def 							_a      (				self      )  ->							Union[str, Any]:
           self.config_tester.run_common_tests()



       def 							_a      (				self      )  ->							Optional[Any]:
           lowercase__       ,       lowercase__		       =						self.model_tester.prepare_config_and_inputs_for_common()

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

               lowercase__		       =						["pixel_values"]
               self.assertListEqual(arg_names[:1]     ,UpperCAmelCase_      )



       def 							_a      (				self      )  ->							int:
           lowercase__       ,       lowercase__		       =						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__      ):
                   lowercase__		       =						self._prepare_for_class(UpperCAmelCase_     ,UpperCAmelCase_      )
                   lowercase__		       =						model_class(UpperCAmelCase_      )

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

                   with self.subTest("JIT Enabled"      ):
                       lowercase__		       =						model_jitted(**UpperCAmelCase_      ).to_tuple()

                   with self.subTest("JIT Disabled"      ):
                       with jax.disable_jit():
                           lowercase__		       =						model_jitted(**UpperCAmelCase_      ).to_tuple()

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



       def 							_a      (				self      )  ->							List[Any]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Union[str, Any]:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase_      )



       def 							_a      (				self      )  ->							str:
           lowercase__		       =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_      )







       @slow
       def 							_a      (				self      )  ->							Tuple:
           for model_class_name in self.all_model_classes:
               lowercase__		       =						model_class_name.from_pretrained("microsoft/beit-base-patch16-224"      )
               lowercase__		       =						model(np.ones((1, 3, 224, 224)      )      )
               self.assertIsNotNone(UpperCAmelCase_      )





def        lowerCamelCase  (				):





    '''simple docstring'''


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



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



       @cached_property
       def 							_a      (				self      )  ->							Optional[int]:
           return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"      ) if is_vision_available() else None



       @slow
       def 							_a      (				self      )  ->							Union[str, Any]:
           lowercase__		       =						FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"      )

           lowercase__		       =						self.default_image_processor
           lowercase__		       =						prepare_img()
           lowercase__		       =						image_processor(images=UpperCAmelCase_     ,return_tensors="np"      ).pixel_values

           # prepare bool_masked_pos
           lowercase__		       =						np.ones((1, 196)     ,dtype=UpperCAmelCase_      )

           # forward pass
           lowercase__		       =						model(pixel_values=UpperCAmelCase_     ,bool_masked_pos=UpperCAmelCase_      )
           lowercase__		       =						outputs.logits

           # verify the logits
           lowercase__		       =						(1, 196, 8_192)
           self.assertEqual(logits.shape     ,UpperCAmelCase_      )

           lowercase__		       =						np.array(
               [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]      )

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



       @slow
       def 							_a      (				self      )  ->							List[str]:
           lowercase__		       =						FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"      )

           lowercase__		       =						self.default_image_processor
           lowercase__		       =						prepare_img()
           lowercase__		       =						image_processor(images=UpperCAmelCase_     ,return_tensors="np"      )

           # forward pass
           lowercase__		       =						model(**UpperCAmelCase_      )
           lowercase__		       =						outputs.logits

           # verify the logits
           lowercase__		       =						(1, 1_000)
           self.assertEqual(logits.shape     ,UpperCAmelCase_      )

           lowercase__		       =						np.array([-1.23_85, -1.09_87, -1.01_08]      )

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

           lowercase__		       =						281
           self.assertEqual(logits.argmax(-1      ).item()     ,UpperCAmelCase_      )







       @slow
       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"      )

           lowercase__		       =						self.default_image_processor
           lowercase__		       =						prepare_img()
           lowercase__		       =						image_processor(images=UpperCAmelCase_     ,return_tensors="np"      )

           # forward pass
           lowercase__		       =						model(**UpperCAmelCase_      )
           lowercase__		       =						outputs.logits

           # verify the logits
           lowercase__		       =						(1, 21_841)
           self.assertEqual(logits.shape     ,UpperCAmelCase_      )

           lowercase__		       =						np.array([1.68_81, -0.27_87, 0.59_01]      )

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

           lowercase__		       =						2_396
           self.assertEqual(logits.argmax(-1      ).item()     ,UpperCAmelCase_      )




