






"""simple docstring"""
from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
			from transformers import (
			    TFBertModel,
			    TFCLIPVisionModel,
			    TFDeiTModel,
			    TFRobertaModel,
			    TFVisionTextDualEncoderModel,
			    TFViTModel,
			    VisionTextDualEncoderConfig,
			)

if is_vision_available():
			from PIL import Image

			from transformers import VisionTextDualEncoderProcessor

def 				snake_case      (   A__						):
			if isinstance(A__	,collections.abc.Iterable						):
						return x
			return (x, x)







@require_tf
class 			UpperCamelCase_       :





	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple	,	lowerCAmelCase_					: Dict	,	lowerCAmelCase_					: List[Any]						)						->			int:
				pass





	def      _SCREAMING_SNAKE_CASE      (      self					: Dict						)						->			Tuple:
				pass





	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple						)						->			Optional[int]:
				pass





	def      _SCREAMING_SNAKE_CASE      (      self					: Dict	,	lowerCAmelCase_					: Optional[int]	,	lowerCAmelCase_					: Any	,	lowerCAmelCase_					: str	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Optional[Any]=None	,	**lowerCAmelCase_					: Union[str, Any]						)						->			Any:
				UpperCAmelCase_     : List[str]			     =					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase_	,	lowerCAmelCase_						)

				UpperCAmelCase_     : Optional[Any]			     =					TFVisionTextDualEncoderModel(lowerCAmelCase_						)

				UpperCAmelCase_     : Tuple			     =					model(input_ids=lowerCAmelCase_	,	pixel_values=lowerCAmelCase_	,	attention_mask=lowerCAmelCase_						)

				self.assertEqual(output["text_embeds"].shape	,	(input_ids.shape[0], config.projection_dim)						)
				self.assertEqual(output["image_embeds"].shape	,	(pixel_values.shape[0], config.projection_dim)						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Any	,	lowerCAmelCase_					: int	,	lowerCAmelCase_					: Optional[Any]	,	lowerCAmelCase_					: int=None	,	**lowerCAmelCase_					: str						)						->			str:
				UpperCAmelCase_     ,	UpperCAmelCase_     : List[str]			     =					self.get_vision_text_model(lowerCAmelCase_	,	lowerCAmelCase_						)
				UpperCAmelCase_     : Union[str, Any]			     =					TFVisionTextDualEncoderModel(vision_model=lowerCAmelCase_	,	text_model=lowerCAmelCase_						)

				UpperCAmelCase_     : Dict			     =					model(input_ids=lowerCAmelCase_	,	pixel_values=lowerCAmelCase_	,	attention_mask=lowerCAmelCase_						)

				self.assertEqual(output["text_embeds"].shape	,	(input_ids.shape[0], model.config.projection_dim)						)
				self.assertEqual(output["image_embeds"].shape	,	(pixel_values.shape[0], model.config.projection_dim)						)





	def      _SCREAMING_SNAKE_CASE      (      self					: List[Any]	,	lowerCAmelCase_					: Union[str, Any]	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Dict	,	lowerCAmelCase_					: int	,	lowerCAmelCase_					: List[Any]=None	,	**lowerCAmelCase_					: str						)						->			Optional[Any]:
				UpperCAmelCase_     ,	UpperCAmelCase_     : Tuple			     =					self.get_vision_text_model(lowerCAmelCase_	,	lowerCAmelCase_						)
				UpperCAmelCase_     : Optional[Any]			     =					{"vision_model": vision_model, "text_model": text_model}
				UpperCAmelCase_     : Optional[int]			     =					TFVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase_						)

				UpperCAmelCase_     : Optional[Any]			     =					model(input_ids=lowerCAmelCase_	,	pixel_values=lowerCAmelCase_	,	attention_mask=lowerCAmelCase_						)

				self.assertEqual(output["text_embeds"].shape	,	(input_ids.shape[0], model.config.projection_dim)						)
				self.assertEqual(output["image_embeds"].shape	,	(pixel_values.shape[0], model.config.projection_dim)						)





	def      _SCREAMING_SNAKE_CASE      (      self					: int	,	lowerCAmelCase_					: Optional[int]	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: List[str]=None	,	**lowerCAmelCase_					: List[str]						)						->			Union[str, Any]:
				UpperCAmelCase_     ,	UpperCAmelCase_     : List[Any]			     =					self.get_vision_text_model(lowerCAmelCase_	,	lowerCAmelCase_						)
				UpperCAmelCase_     : Optional[int]			     =					TFVisionTextDualEncoderModel(vision_model=lowerCAmelCase_	,	text_model=lowerCAmelCase_						)

				UpperCAmelCase_     : Optional[Any]			     =					model(input_ids=lowerCAmelCase_	,	pixel_values=lowerCAmelCase_	,	attention_mask=lowerCAmelCase_						)
				UpperCAmelCase_     : Any			     =					output[0].numpy()

				with tempfile.TemporaryDirectory() as tmpdirname:
							model.save_pretrained(lowerCAmelCase_						)
							UpperCAmelCase_     : Tuple			     =					TFVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase_						)

							UpperCAmelCase_     : List[Any]			     =					model(input_ids=lowerCAmelCase_	,	pixel_values=lowerCAmelCase_	,	attention_mask=lowerCAmelCase_						)
							UpperCAmelCase_     : List[str]			     =					after_output[0].numpy()
							UpperCAmelCase_     : Any			     =					np.amax(np.abs(out_a - out_a						)						)
							self.assertLessEqual(lowerCAmelCase_	,	1e-5						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]	,	lowerCAmelCase_					: Optional[Any]	,	lowerCAmelCase_					: str	,	lowerCAmelCase_					: int	,	lowerCAmelCase_					: Any	,	lowerCAmelCase_					: Optional[Any]=None	,	**lowerCAmelCase_					: int						)						->			Union[str, Any]:
				UpperCAmelCase_     ,	UpperCAmelCase_     : Dict			     =					self.get_vision_text_model(lowerCAmelCase_	,	lowerCAmelCase_						)
				UpperCAmelCase_     : Tuple			     =					TFVisionTextDualEncoderModel(vision_model=lowerCAmelCase_	,	text_model=lowerCAmelCase_						)

				UpperCAmelCase_     : List[str]			     =					model(
				    input_ids=lowerCAmelCase_	,	pixel_values=lowerCAmelCase_	,	attention_mask=lowerCAmelCase_	,	output_attentions=lowerCAmelCase_						)

				UpperCAmelCase_     : str			     =					output.vision_model_output.attentions
				self.assertEqual(len(lowerCAmelCase_						)	,	vision_config.num_hidden_layers						)

				# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
				UpperCAmelCase_     : Dict			     =					to_atuple(vision_model.config.image_size						)
				UpperCAmelCase_     : Optional[Any]			     =					to_atuple(vision_model.config.patch_size						)
				UpperCAmelCase_     : Union[str, Any]			     =					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
				UpperCAmelCase_     : Dict			     =					num_patches + 1
				self.assertEqual(vision_attentions[0].shape[-3:]	,	(vision_config.num_attention_heads, seq_len, seq_len)						)

				UpperCAmelCase_     : Tuple			     =					output.text_model_output.attentions
				self.assertEqual(len(lowerCAmelCase_						)	,	text_config.num_hidden_layers						)

				self.assertEqual(
				    text_attentions[0].shape[-3:]	,	(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])	,	)





	def      _SCREAMING_SNAKE_CASE      (      self					: int	,	lowerCAmelCase_					: np.ndarray	,	lowerCAmelCase_					: np.ndarray	,	lowerCAmelCase_					: float						)						->			str:
				UpperCAmelCase_     : Dict			     =					np.abs((a - b)						).max()
				self.assertLessEqual(lowerCAmelCase_	,	lowerCAmelCase_	,	f"""Difference between torch and flax is {diff} (>= {tol})."""						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Any						)						->			List[str]:
				UpperCAmelCase_     : List[Any]			     =					self.prepare_config_and_inputs()
				self.check_vision_text_dual_encoder_model(**lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: List[str]						)						->			List[str]:
				UpperCAmelCase_     : int			     =					self.prepare_config_and_inputs()
				self.check_model_from_pretrained_configs(**lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: int						)						->			List[str]:
				UpperCAmelCase_     : Tuple			     =					self.prepare_config_and_inputs()
				self.check_vision_text_dual_encoder_from_pretrained(**lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: Union[str, Any]						)						->			Tuple:
				UpperCAmelCase_     : str			     =					self.prepare_config_and_inputs()
				self.check_save_load(**lowerCAmelCase_						)





	def      _SCREAMING_SNAKE_CASE      (      self					: int						)						->			List[str]:
				UpperCAmelCase_     : Tuple			     =					self.prepare_config_and_inputs()
				self.check_vision_text_output_attention(**lowerCAmelCase_						)





	@slow
	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple						)						->			List[Any]:
				UpperCAmelCase_     ,	UpperCAmelCase_     : Tuple			     =					self.get_pretrained_model_and_inputs()

				UpperCAmelCase_     : str			     =					model_a(**lowerCAmelCase_						)
				UpperCAmelCase_     : Optional[Any]			     =					outputs[0].numpy()

				with tempfile.TemporaryDirectory() as tmp_dirname:
							model_a.save_pretrained(lowerCAmelCase_						)
							UpperCAmelCase_     : Optional[int]			     =					TFVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase_						)

							UpperCAmelCase_     : str			     =					model_a(**lowerCAmelCase_						)
							UpperCAmelCase_     : str			     =					after_outputs[0].numpy()
							UpperCAmelCase_     : Optional[int]			     =					np.amax(np.abs(out_a - out_a						)						)
							self.assertLessEqual(lowerCAmelCase_	,	1e-5						)







@require_tf
class 			UpperCamelCase_       (__A				,						unittest.TestCase     ):





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]						)						->			int:
				UpperCAmelCase_     : List[str]			     =					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
				    "hf-internal-testing/tiny-random-vit"	,	"hf-internal-testing/tiny-random-bert"						)
				UpperCAmelCase_     : Optional[int]			     =					13
				UpperCAmelCase_     : Optional[Any]			     =					floats_tensor(
				    [
				        batch_size,
				        model.vision_model.config.num_channels,
				        model.vision_model.config.image_size,
				        model.vision_model.config.image_size,
				    ]						)
				UpperCAmelCase_     : int			     =					ids_tensor([batch_size, 4]	,	model.text_model.config.vocab_size						)
				UpperCAmelCase_     : Tuple			     =					random_attention_mask([batch_size, 4]						)
				UpperCAmelCase_     : Optional[int]			     =					{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

				return model, inputs





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]	,	lowerCAmelCase_					: Dict	,	lowerCAmelCase_					: Optional[Any]						)						->			List[Any]:
				UpperCAmelCase_     : int			     =					TFViTModel(lowerCAmelCase_	,	name="vision_model"						)
				UpperCAmelCase_     : Optional[Any]			     =					TFBertModel(lowerCAmelCase_	,	name="text_model"						)
				return vision_model, text_model





	def      _SCREAMING_SNAKE_CASE      (      self					: List[str]						)						->			Optional[int]:
				UpperCAmelCase_     : List[Any]			     =					TFViTModelTester(self						)
				UpperCAmelCase_     : List[Any]			     =					TFBertModelTester(self						)
				UpperCAmelCase_     : Optional[int]			     =					vit_model_tester.prepare_config_and_inputs()
				UpperCAmelCase_     : Any			     =					bert_model_tester.prepare_config_and_inputs()

				UpperCAmelCase_     ,	UpperCAmelCase_     ,	UpperCAmelCase_     : Tuple			     =					vision_config_and_inputs

				(
				    (
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	
				)     : List[Any]			     =					text_config_and_inputs

				return {
				    "text_config": text_config,
				    "vision_config": vision_config,
				    "pixel_values": pixel_values,
				    "attention_mask": input_mask,
				    "input_ids": input_ids,
				    "text_token_type_ids": token_type_ids,
				    "text_sequence_labels": sequence_labels,
				    "text_token_labels": token_labels,
				    "text_choice_labels": choice_labels,
				}







@require_tf
class 			UpperCamelCase_       (__A				,						unittest.TestCase     ):





	def      _SCREAMING_SNAKE_CASE      (      self					: int						)						->			Optional[int]:
				# DeiT repo doesn't have TF weights, but we don't actually use the weights at all so let's
				# just reinitialize it.
				UpperCAmelCase_     : Tuple			     =					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
				    "Rocketknight1/tiny-random-deit-tf"	,	"hf-internal-testing/tiny-random-roberta"						)
				UpperCAmelCase_     : Any			     =					13
				UpperCAmelCase_     : List[str]			     =					floats_tensor(
				    [
				        batch_size,
				        model.vision_model.config.num_channels,
				        model.vision_model.config.image_size,
				        model.vision_model.config.image_size,
				    ]						)
				UpperCAmelCase_     : Optional[Any]			     =					ids_tensor([batch_size, 4]	,	model.text_model.config.vocab_size						)
				UpperCAmelCase_     : Union[str, Any]			     =					random_attention_mask([batch_size, 4]						)
				UpperCAmelCase_     : List[str]			     =					{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

				return model, inputs





	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple	,	lowerCAmelCase_					: Optional[int]	,	lowerCAmelCase_					: Tuple	,	lowerCAmelCase_					: List[str]	,	lowerCAmelCase_					: Optional[int]	,	lowerCAmelCase_					: int=None	,	**lowerCAmelCase_					: Optional[Any]						)						->			Union[str, Any]:
				UpperCAmelCase_     ,	UpperCAmelCase_     : Dict			     =					self.get_vision_text_model(lowerCAmelCase_	,	lowerCAmelCase_						)
				UpperCAmelCase_     : Union[str, Any]			     =					TFVisionTextDualEncoderModel(vision_model=lowerCAmelCase_	,	text_model=lowerCAmelCase_						)

				UpperCAmelCase_     : List[Any]			     =					model(
				    input_ids=lowerCAmelCase_	,	pixel_values=lowerCAmelCase_	,	attention_mask=lowerCAmelCase_	,	output_attentions=lowerCAmelCase_						)

				UpperCAmelCase_     : Tuple			     =					output.vision_model_output.attentions
				self.assertEqual(len(lowerCAmelCase_						)	,	vision_config.num_hidden_layers						)

				# in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
				UpperCAmelCase_     : List[Any]			     =					to_atuple(vision_model.config.image_size						)
				UpperCAmelCase_     : Union[str, Any]			     =					to_atuple(vision_model.config.patch_size						)
				UpperCAmelCase_     : Optional[int]			     =					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
				UpperCAmelCase_     : str			     =					num_patches + 2
				self.assertEqual(vision_attentions[0].shape[-3:]	,	(vision_config.num_attention_heads, seq_len, seq_len)						)

				UpperCAmelCase_     : List[str]			     =					output.text_model_output.attentions
				self.assertEqual(len(lowerCAmelCase_						)	,	text_config.num_hidden_layers						)

				self.assertEqual(
				    text_attentions[0].shape[-3:]	,	(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])	,	)





	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple	,	lowerCAmelCase_					: Any	,	lowerCAmelCase_					: Tuple						)						->			Union[str, Any]:
				UpperCAmelCase_     : Optional[Any]			     =					TFDeiTModel(lowerCAmelCase_	,	name="vision_model"						)
				UpperCAmelCase_     : str			     =					TFRobertaModel(lowerCAmelCase_	,	name="text_model"						)
				return vision_model, text_model





	def      _SCREAMING_SNAKE_CASE      (      self					: int						)						->			str:
				UpperCAmelCase_     : Dict			     =					TFDeiTModelTester(self						)
				UpperCAmelCase_     : Optional[Any]			     =					TFRobertaModelTester(self						)
				UpperCAmelCase_     : Union[str, Any]			     =					vit_model_tester.prepare_config_and_inputs()
				UpperCAmelCase_     : Tuple			     =					bert_model_tester.prepare_config_and_inputs()

				UpperCAmelCase_     ,	UpperCAmelCase_     ,	UpperCAmelCase_     : Dict			     =					vision_config_and_inputs

				(
				    (
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	
				)     : Optional[int]			     =					text_config_and_inputs

				return {
				    "text_config": text_config,
				    "vision_config": vision_config,
				    "pixel_values": pixel_values,
				    "attention_mask": input_mask,
				    "input_ids": input_ids,
				    "text_token_type_ids": token_type_ids,
				    "text_sequence_labels": sequence_labels,
				    "text_token_labels": token_labels,
				    "text_choice_labels": choice_labels,
				}







@require_tf
class 			UpperCamelCase_       (__A				,						unittest.TestCase     ):





	def      _SCREAMING_SNAKE_CASE      (      self					: Dict						)						->			Any:
				UpperCAmelCase_     : Optional[int]			     =					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
				    "Rocketknight1/tiny-random-clip-tf"	,	"hf-internal-testing/tiny-random-bert"						)
				UpperCAmelCase_     : List[Any]			     =					13
				UpperCAmelCase_     : Dict			     =					floats_tensor(
				    [
				        batch_size,
				        model.vision_model.config.num_channels,
				        model.vision_model.config.image_size,
				        model.vision_model.config.image_size,
				    ]						)
				UpperCAmelCase_     : Optional[int]			     =					ids_tensor([batch_size, 4]	,	model.text_model.config.vocab_size						)
				UpperCAmelCase_     : Optional[Any]			     =					random_attention_mask([batch_size, 4]						)
				UpperCAmelCase_     : Dict			     =					{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

				return model, inputs





	def      _SCREAMING_SNAKE_CASE      (      self					: Any	,	lowerCAmelCase_					: Dict	,	lowerCAmelCase_					: Optional[int]						)						->			Optional[int]:
				UpperCAmelCase_     : Tuple			     =					TFCLIPVisionModel(lowerCAmelCase_	,	name="vision_model"						)
				UpperCAmelCase_     : Optional[int]			     =					TFBertModel(lowerCAmelCase_	,	name="text_model"						)
				return vision_model, text_model





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]						)						->			Union[str, Any]:
				UpperCAmelCase_     : str			     =					TFCLIPVisionModelTester(self						)
				UpperCAmelCase_     : str			     =					TFBertModelTester(self						)
				UpperCAmelCase_     : str			     =					clip_model_tester.prepare_config_and_inputs()
				UpperCAmelCase_     : int			     =					bert_model_tester.prepare_config_and_inputs()

				UpperCAmelCase_     ,	UpperCAmelCase_     : List[str]			     =					vision_config_and_inputs

				(
				    (
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	(
				    UpperCAmelCase_
				)     ,	
				)     : Optional[int]			     =					text_config_and_inputs

				return {
				    "text_config": text_config,
				    "vision_config": vision_config,
				    "pixel_values": pixel_values,
				    "attention_mask": input_mask,
				    "input_ids": input_ids,
				    "text_token_type_ids": token_type_ids,
				    "text_sequence_labels": sequence_labels,
				    "text_token_labels": token_labels,
				    "text_choice_labels": choice_labels,
				}







@require_vision
@require_tf
class 			UpperCamelCase_       (unittest.TestCase     ):





	@slow
	def      _SCREAMING_SNAKE_CASE      (      self					: str						)						->			int:
				UpperCAmelCase_     : Any			     =					TFVisionTextDualEncoderModel.from_pretrained(
				    "clip-italian/clip-italian"	,	logit_scale_init_value=1.0	,	from_pt=lowerCAmelCase_						)
				UpperCAmelCase_     : Tuple			     =					VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"						)

				UpperCAmelCase_     : str			     =					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
				UpperCAmelCase_     : Dict			     =					processor(
				    text=["una foto di un gatto", "una foto di un cane"]	,	images=lowerCAmelCase_	,	padding=lowerCAmelCase_	,	return_tensors="np"						)

				UpperCAmelCase_     : Any			     =					model(**lowerCAmelCase_						)

				# verify the logits
				self.assertEqual(outputs.logits_per_image.shape	,	(inputs.pixel_values.shape[0], inputs.input_ids.shape[0])						)
				self.assertEqual(
				    outputs.logits_per_text.shape	,	(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])	,	)

				UpperCAmelCase_     : List[Any]			     =					np.array([[1.2_2_8_4_7_2_7, 0.3_1_0_4_1_2_2]]						)

				self.assertTrue(np.allclose(outputs.logits_per_image.numpy()	,	lowerCAmelCase_	,	atol=1e-3						)						)











"""simple docstring"""
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_tf_available():
			import tensorflow as tf

			from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
			from ..tf_utils import stable_softmax

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

lowerCamelCase_					=      logging.get_logger(__name__)







@add_end_docstrings(__A     )
class 			UpperCamelCase_       (__A     ):





	def __init__(      self					: int	,	*lowerCAmelCase_					: Tuple	,	**lowerCAmelCase_					: List[str]						)						->			Optional[Any]:
				super().__init__(*lowerCAmelCase_	,	**lowerCAmelCase_						)
				requires_backends(self	,	"vision"						)
				self.check_model_type(
				    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
				    if self.framework == "tf"
				    else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING						)





	def      _SCREAMING_SNAKE_CASE      (      self					: str	,	lowerCAmelCase_					: Optional[int]=None						)						->			List[Any]:
				UpperCAmelCase_     : str			     =					{}
				if top_k is not None:
							UpperCAmelCase_     : List[str]			     =					top_k
				return {}, {}, postprocess_params





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





	def      _SCREAMING_SNAKE_CASE      (      self					: str	,	lowerCAmelCase_					: str						)						->			Any:
				UpperCAmelCase_     : Tuple			     =					load_image(lowerCAmelCase_						)
				UpperCAmelCase_     : Dict			     =					self.image_processor(images=lowerCAmelCase_	,	return_tensors=self.framework						)
				return model_inputs





	def      _SCREAMING_SNAKE_CASE      (      self					: List[Any]	,	lowerCAmelCase_					: Dict						)						->			str:
				UpperCAmelCase_     : Any			     =					self.model(**lowerCAmelCase_						)
				return model_outputs





	def      _SCREAMING_SNAKE_CASE      (      self					: Tuple	,	lowerCAmelCase_					: Union[str, Any]	,	lowerCAmelCase_					: Optional[int]=5						)						->			Any:
				if top_k > self.model.config.num_labels:
							UpperCAmelCase_     : int			     =					self.model.config.num_labels

				if self.framework == "pt":
							UpperCAmelCase_     : str			     =					model_outputs.logits.softmax(-1						)[0]
							UpperCAmelCase_     ,	UpperCAmelCase_     : Tuple			     =					probs.topk(lowerCAmelCase_						)
				elif self.framework == "tf":
							UpperCAmelCase_     : str			     =					stable_softmax(model_outputs.logits	,	axis=-1						)[0]
							UpperCAmelCase_     : Union[str, Any]			     =					tf.math.top_k(lowerCAmelCase_	,	k=lowerCAmelCase_						)
							UpperCAmelCase_     ,	UpperCAmelCase_     : List[Any]			     =					topk.values.numpy(), topk.indices.numpy()
				else:
							raise ValueError(f"""Unsupported framework: {self.framework}"""						)

				UpperCAmelCase_     : int			     =					scores.tolist()
				UpperCAmelCase_     : Optional[Any]			     =					ids.tolist()
				return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(lowerCAmelCase_	,	lowerCAmelCase_						)]



