

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







   '''simple docstring'''




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

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





   def      _snake_case		(       self					:			List[Any]			):
       SCREAMING_SNAKE_CASE										=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

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

       SCREAMING_SNAKE_CASE										=   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=_lowercase     , initializer_range=self.initializer_range     , )

       return config, pixel_values, labels





   def      _snake_case		(       self					:			Union[str, Any]     , __lowerCamelCase					:			int     , __lowerCamelCase					:			int     , __lowerCamelCase					:			Any			):
       SCREAMING_SNAKE_CASE										=   FlaxBeitModel(config=_lowercase			)
       SCREAMING_SNAKE_CASE										=   model(_lowercase			)
       self.parent.assertEqual(result.last_hidden_state.shape     , (self.batch_size, self.seq_length, self.hidden_size)			)





   def      _snake_case		(       self					:			int     , __lowerCamelCase					:			List[Any]     , __lowerCamelCase					:			Tuple     , __lowerCamelCase					:			Union[str, Any]			):
       SCREAMING_SNAKE_CASE										=   FlaxBeitForMaskedImageModeling(config=_lowercase			)
       SCREAMING_SNAKE_CASE										=   model(_lowercase			)
       self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.seq_length - 1, self.vocab_size)			)





   def      _snake_case		(       self					:			List[str]     , __lowerCamelCase					:			str     , __lowerCamelCase					:			List[Any]     , __lowerCamelCase					:			Optional[Any]			):
       SCREAMING_SNAKE_CASE										=   self.type_sequence_label_size
       SCREAMING_SNAKE_CASE										=   FlaxBeitForImageClassification(config=_lowercase			)
       SCREAMING_SNAKE_CASE										=   model(_lowercase			)
       self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.type_sequence_label_size)			)

       # test greyscale images
       SCREAMING_SNAKE_CASE										=   1
       SCREAMING_SNAKE_CASE										=   FlaxBeitForImageClassification(_lowercase			)

       SCREAMING_SNAKE_CASE										=   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
       SCREAMING_SNAKE_CASE										=   model(_lowercase			)





   def      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE										=   self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE 
       )							,	(
           SCREAMING_SNAKE_CASE 
       )							,	(
           SCREAMING_SNAKE_CASE 
       )							,	 
       )										=   config_and_inputs
       SCREAMING_SNAKE_CASE										=   {"pixel_values": pixel_values}
       return config, inputs_dict





@require_flax
class 			_SCREAMING_SNAKE_CASE (   __snake_case					,       unittest.TestCase		):







   '''simple docstring'''




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





   def      _snake_case		(       self					:			Dict			):
       SCREAMING_SNAKE_CASE										=   FlaxBeitModelTester(self			)
       SCREAMING_SNAKE_CASE										=   ConfigTester(self     , config_class=_lowercase     , has_text_modality=_lowercase     , hidden_size=37			)





   def      _snake_case		(       self					:			Optional[Any]			):
       self.config_tester.run_common_tests()





   def      _snake_case		(       self					:			List[Any]			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs_for_common()

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

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





   def      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__			):
               SCREAMING_SNAKE_CASE										=   self._prepare_for_class(_lowercase     , _lowercase			)
               SCREAMING_SNAKE_CASE										=   model_class(_lowercase			)

               @jax.jit
               def model_jitted(__lowerCamelCase					:			List[Any]     , **__lowerCamelCase					:			Tuple			):
                   return model(pixel_values=_lowercase     , **_lowercase			)

               with self.subTest("JIT Enabled"			):
                   SCREAMING_SNAKE_CASE										=   model_jitted(**_lowercase			).to_tuple()

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

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





   def      _snake_case		(       self					:			Tuple			):
       SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_lowercase			)





   def      _snake_case		(       self					:			Optional[int]			):
       SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*_lowercase			)





   def      _snake_case		(       self					:			Any			):
       SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_lowercase			)





   @slow
   def      _snake_case		(       self					:			str			):
       for model_class_name in self.all_model_classes:
           SCREAMING_SNAKE_CASE										=   model_class_name.from_pretrained("microsoft/beit-base-patch16-224"			)
           SCREAMING_SNAKE_CASE										=   model(np.ones((1, 3, 224, 224)			)			)
           self.assertIsNotNone(_lowercase			)





def 	__a     (      ):
    SCREAMING_SNAKE_CASE										=   Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
    return image





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







   '''simple docstring'''




   @cached_property
   def      _snake_case		(       self					:			List[Any]			):
       return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"			) if is_vision_available() else None





   @slow
   def      _snake_case		(       self					:			Optional[Any]			):
       SCREAMING_SNAKE_CASE										=   FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"			)

       SCREAMING_SNAKE_CASE										=   self.default_image_processor
       SCREAMING_SNAKE_CASE										=   prepare_img()
       SCREAMING_SNAKE_CASE										=   image_processor(images=_lowercase     , return_tensors="np"			).pixel_values

       # prepare bool_masked_pos
       SCREAMING_SNAKE_CASE										=   np.ones((1, 196)     , dtype=_lowercase			)

       # forward pass
       SCREAMING_SNAKE_CASE										=   model(pixel_values=_lowercase     , bool_masked_pos=_lowercase			)
       SCREAMING_SNAKE_CASE										=   outputs.logits

       # verify the logits
       SCREAMING_SNAKE_CASE										=   (1, 196, 8192)
       self.assertEqual(logits.shape     , _lowercase			)

       SCREAMING_SNAKE_CASE										=   np.array(
           [[-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]]			)

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





   @slow
   def      _snake_case		(       self					:			Any			):
       SCREAMING_SNAKE_CASE										=   FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"			)

       SCREAMING_SNAKE_CASE										=   self.default_image_processor
       SCREAMING_SNAKE_CASE										=   prepare_img()
       SCREAMING_SNAKE_CASE										=   image_processor(images=_lowercase     , return_tensors="np"			)

       # forward pass
       SCREAMING_SNAKE_CASE										=   model(**_lowercase			)
       SCREAMING_SNAKE_CASE										=   outputs.logits

       # verify the logits
       SCREAMING_SNAKE_CASE										=   (1, 1000)
       self.assertEqual(logits.shape     , _lowercase			)

       SCREAMING_SNAKE_CASE										=   np.array([-1.2_385, -1.0_987, -1.0_108]			)

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

       SCREAMING_SNAKE_CASE										=   281
       self.assertEqual(logits.argmax(-1			).item()     , _lowercase			)





   @slow
   def      _snake_case		(       self					:			Dict			):
       SCREAMING_SNAKE_CASE										=   FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"			)

       SCREAMING_SNAKE_CASE										=   self.default_image_processor
       SCREAMING_SNAKE_CASE										=   prepare_img()
       SCREAMING_SNAKE_CASE										=   image_processor(images=_lowercase     , return_tensors="np"			)

       # forward pass
       SCREAMING_SNAKE_CASE										=   model(**_lowercase			)
       SCREAMING_SNAKE_CASE										=   outputs.logits

       # verify the logits
       SCREAMING_SNAKE_CASE										=   (1, 21841)
       self.assertEqual(logits.shape     , _lowercase			)

       SCREAMING_SNAKE_CASE										=   np.array([1.6_881, -0.2_787, 0.5_901]			)

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

       SCREAMING_SNAKE_CASE										=   2396
       self.assertEqual(logits.argmax(-1			).item()     , _lowercase			)





'''simple docstring'''
from collections import UserDict
from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
							from PIL import Image

							from ..image_utils import load_image

if is_torch_available():
							from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING

if is_tf_available():
							from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
							from ..tf_utils import stable_softmax

_SCREAMING_SNAKE_CASE			=						logging.get_logger(__name__)
@add_end_docstrings(lowerCAmelCase				)
class   __UpperCAmelCase				(     lowerCAmelCase				):



					'''simple docstring'''





					def __init__(	self     :						int      ,       **_lowercase     :						List[Any])    ->							Union[str, Any]:
								super().__init__(**_lowercase)

								requires_backends(self      ,       'vision')
								self.check_model_type(
								    TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
								    if self.framework == 'tf'
								    else MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING)



					def __call__(	self     :						Optional[Any]      ,       _lowercase     :						Union[str, List[str], "Image", List["Image"]]      ,       **_lowercase     :						List[Any])    ->							Any:
								return super().__call__(_lowercase      ,       **_lowercase)



					def     __snake_case						(	self     :						int      ,       **_lowercase     :						Union[str, Any])    ->							Any:
								A_									= {}
								if "candidate_labels" in kwargs:
											A_									= kwargs['candidate_labels']
								if "hypothesis_template" in kwargs:
											A_									= kwargs['hypothesis_template']

								return preprocess_params, {}, {}



					def     __snake_case						(	self     :						List[str]      ,       _lowercase     :						Tuple      ,       _lowercase     :						Any=None      ,       _lowercase     :						Optional[int]="This is a photo of {}.")    ->							Union[str, Any]:
								A_									= load_image(_lowercase)
								A_									= self.image_processor(images=[image]      ,       return_tensors=self.framework)
								A_									= candidate_labels
								A_									= [hypothesis_template.format(_lowercase) for x in candidate_labels]
								A_									= self.tokenizer(_lowercase      ,       return_tensors=self.framework      ,       padding=_lowercase)
								A_									= [text_inputs]
								return inputs



					def     __snake_case						(	self     :						Optional[int]      ,       _lowercase     :						Tuple)    ->							Optional[int]:
								A_									= model_inputs.pop('candidate_labels')
								A_									= model_inputs.pop('text_inputs')
								if isinstance(text_inputs[0]      ,       _lowercase):
											A_									= text_inputs[0]
								else:
											# Batching case.
											A_									= text_inputs[0][0]

								A_									= self.model(**_lowercase      ,       **_lowercase)

								A_									= {
								    'candidate_labels': candidate_labels,
								    'logits': outputs.logits_per_image,
								}
								return model_outputs



					def     __snake_case						(	self     :						List[str]      ,       _lowercase     :						int)    ->							Optional[int]:
								A_									= model_outputs.pop('candidate_labels')
								A_									= model_outputs['logits'][0]
								if self.framework == "pt":
											A_									= logits.softmax(dim=-1).squeeze(-1)
											A_									= probs.tolist()
											if not isinstance(_lowercase      ,       _lowercase):
														A_									= [scores]
								elif self.framework == "tf":
											A_									= stable_softmax(_lowercase      ,       axis=-1)
											A_									= probs.numpy().tolist()
								else:
											raise ValueError(F'Unsupported framework: {self.framework}')

								A_									= [
								    {'score': score, 'label': candidate_label}
								    for score, candidate_label in sorted(zip(_lowercase      ,       _lowercase)      ,       key=lambda _lowercase: -x[0])
								]
								return result




