




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



					'''simple docstring'''





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

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



					def     __snake_case						(	self     :						List[Any])    ->							int:
								A_									= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

								A_									= 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]      ,       _lowercase     :						int      ,       _lowercase     :						int      ,       _lowercase     :						Any)    ->							Optional[Any]:
								A_									= FlaxBeitModel(config=_lowercase)
								A_									= model(_lowercase)
								self.parent.assertEqual(result.last_hidden_state.shape      ,       (self.batch_size, self.seq_length, self.hidden_size))



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



					def     __snake_case						(	self     :						List[str]      ,       _lowercase     :						str      ,       _lowercase     :						List[Any]      ,       _lowercase     :						Optional[Any])    ->							Optional[Any]:
								A_									= self.type_sequence_label_size
								A_									= FlaxBeitForImageClassification(config=_lowercase)
								A_									= model(_lowercase)
								self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.type_sequence_label_size))

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

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



					def     __snake_case						(	self     :						List[str])    ->							Dict:
								A_									= self.prepare_config_and_inputs()
								(
								    (
								    A_
								)		, (
								    A_
								)		, (
								    A_
								)		, 
								)									= config_and_inputs
								A_									= {'pixel_values': pixel_values}
								return config, inputs_dict





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



					'''simple docstring'''





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



					def     __snake_case						(	self     :						Dict)    ->							None:
								A_									= FlaxBeitModelTester(self)
								A_									= ConfigTester(self      ,       config_class=_lowercase      ,       has_text_modality=_lowercase      ,       hidden_size=37)



					def     __snake_case						(	self     :						Optional[Any])    ->							Optional[int]:
								self.config_tester.run_common_tests()



					def     __snake_case						(	self     :						List[Any])    ->							List[Any]:
								A_		, A_									= self.model_tester.prepare_config_and_inputs_for_common()

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

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



					def     __snake_case						(	self     :						List[str])    ->							Any:
								A_		, A_									= self.model_tester.prepare_config_and_inputs_for_common()

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

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

														with self.subTest('JIT Enabled'):
																	A_									= model_jitted(**_lowercase).to_tuple()

														with self.subTest('JIT Disabled'):
																	with jax.disable_jit():
																				A_									= 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)    ->							Any:
								A_									= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_lowercase)



					def     __snake_case						(	self     :						Optional[int])    ->							Any:
								A_									= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*_lowercase)



					def     __snake_case						(	self     :						Any)    ->							Dict:
								A_									= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_lowercase)



					@slow
					def     __snake_case						(	self     :						str)    ->							Union[str, Any]:
								for model_class_name in self.all_model_classes:
											A_									= model_class_name.from_pretrained('microsoft/beit-base-patch16-224')
											A_									= model(np.ones((1, 3, 224, 224)))
											self.assertIsNotNone(_lowercase)





def 			lowerCamelCase(							)		->		Union[str, Any]:
			A_									= Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'     )
			return image






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



					'''simple docstring'''





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



					@slow
					def     __snake_case						(	self     :						Optional[Any])    ->							int:
								A_									= FlaxBeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k')

								A_									= self.default_image_processor
								A_									= prepare_img()
								A_									= image_processor(images=_lowercase      ,       return_tensors='np').pixel_values

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

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

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

								A_									= 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]      ,       _lowercase      ,       atol=1E-2))



					@slow
					def     __snake_case						(	self     :						Any)    ->							str:
								A_									= FlaxBeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224')

								A_									= self.default_image_processor
								A_									= prepare_img()
								A_									= image_processor(images=_lowercase      ,       return_tensors='np')

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

								# verify the logits
								A_									= (1, 1_000)
								self.assertEqual(logits.shape      ,       _lowercase)

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

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

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



					@slow
					def     __snake_case						(	self     :						Dict)    ->							Union[str, Any]:
								A_									= FlaxBeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k')

								A_									= self.default_image_processor
								A_									= prepare_img()
								A_									= image_processor(images=_lowercase      ,       return_tensors='np')

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

								# verify the logits
								A_									= (1, 21_841)
								self.assertEqual(logits.shape      ,       _lowercase)

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

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

								A_									= 2_396
								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




