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

_UpperCAmelCase		:       Tuple						    =	logging.get_logger(__name__)







@add_end_docstrings(lowercase_      )
class 			lowercase    (	lowercase_      ):




				def __init__( self				,       **snake_case	):
											super().__init__(**__A	)

											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				,       snake_case				,       **snake_case	):
											return super().__call__(__A				,       **__A	)




				def    a							( self				,       **snake_case	):
											snake_case_          =   {}
											if "candidate_labels" in kwargs:
																		snake_case_          =   kwargs['candidate_labels']
											if "hypothesis_template" in kwargs:
																		snake_case_          =   kwargs['hypothesis_template']

											return preprocess_params, {}, {}




				def    a							( self				,       snake_case				,       snake_case=None				,       snake_case="This is a photo of {}."	):
											snake_case_          =   load_image(__A	)
											snake_case_          =   self.image_processor(images=[image]				,       return_tensors=self.framework	)
											snake_case_          =   candidate_labels
											snake_case_          =   [hypothesis_template.format(__A	) for x in candidate_labels]
											snake_case_          =   self.tokenizer(__A				,       return_tensors=self.framework				,       padding=__A	)
											snake_case_          =   [text_inputs]
											return inputs




				def    a							( self				,       snake_case	):
											snake_case_          =   model_inputs.pop('candidate_labels'	)
											snake_case_          =   model_inputs.pop('text_inputs'	)
											if isinstance(text_inputs[0]				,       __A	):
																		snake_case_          =   text_inputs[0]
											else:
																		# Batching case.
																		snake_case_          =   text_inputs[0][0]

											snake_case_          =   self.model(**__A				,       **__A	)

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




				def    a							( self				,       snake_case	):
											snake_case_          =   model_outputs.pop('candidate_labels'	)
											snake_case_          =   model_outputs['logits'][0]
											if self.framework == "pt":
																		snake_case_          =   logits.softmax(dim=-1	).squeeze(-1	)
																		snake_case_          =   probs.tolist()
																		if not isinstance(__A				,       __A	):
																									snake_case_          =   [scores]
											elif self.framework == "tf":
																		snake_case_          =   stable_softmax(__A				,       axis=-1	)
																		snake_case_          =   probs.numpy().tolist()
											else:
																		raise ValueError(F'''Unsupported framework: {self.framework}'''	)

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







import unittest

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device


if is_torch_available():
					import torch

					from transformers import AutoModelForImageClassification

if is_vision_available():
					from transformers import AutoImageProcessor







@require_torch
@require_vision
class 			lowercase    (	unittest.TestCase      ):




				@slow
				def    a							( self	):
											snake_case_          =   AutoImageProcessor.from_pretrained('microsoft/dit-base-finetuned-rvlcdip'	)
											snake_case_          =   AutoModelForImageClassification.from_pretrained('microsoft/dit-base-finetuned-rvlcdip'	)
											model.to(snake_case	)

											from datasets import load_dataset

											snake_case_          =   load_dataset('nielsr/rvlcdip-demo'	)

											snake_case_          =   dataset['train'][0]['image'].convert('RGB'	)

											snake_case_          =   image_processor(snake_case				,       return_tensors='pt'	).to(snake_case	)

											# forward pass
											with torch.no_grad():
																		snake_case_          =   model(**snake_case	)
																		snake_case_          =   outputs.logits

											snake_case_          =   torch.Size((1, 16)	)
											self.assertEqual(logits.shape				,       snake_case	)

											snake_case_          =   torch.tensor(
											    [-0.41_58, -0.40_92, -0.43_47]				,       device=snake_case				,       dtype=torch.float				,       )
											self.assertTrue(torch.allclose(logits[0, :3]				,       snake_case				,       atol=1e-4	)	)




