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_								=							logging.get_logger(__name__)





@add_end_docstrings(_a			)
class      lowerCAmelCase						(					_a			):



     def __init__(					self				,      **lowerCAmelCase__       ):
            super().__init__(**lowerCAmelCase__       )

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



     def 			a__  (					self				,      **lowerCAmelCase__       ):
            _A=	{}
            if "candidate_labels" in kwargs:
                   _A=	kwargs['candidate_labels']
            if "hypothesis_template" in kwargs:
                   _A=	kwargs['hypothesis_template']

            return preprocess_params, {}, {}



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



     def 			a__  (					self				,      lowerCAmelCase__       ):
            _A=	model_inputs.pop('candidate_labels'       )
            _A=	model_inputs.pop('text_inputs'       )
            if isinstance(text_inputs[0]				,      lowerCAmelCase__       ):
                   _A=	text_inputs[0]
            else:
                   # Batching case.
                   _A=	text_inputs[0][0]

            _A=	self.model(**lowerCAmelCase__				,      **lowerCAmelCase__       )

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





     def 			a__  (					self				,      lowerCAmelCase__       ):
            _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(lowerCAmelCase__				,      lowerCAmelCase__       ):
                          _A=	[scores]
            elif self.framework == "tf":
                   _A=	stable_softmax(lowerCAmelCase__				,      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(lowerCAmelCase__				,      lowerCAmelCase__       )				,      key=lambda lowerCAmelCase__       : -x[0]       )
            ]
            return result
import pickle
import unittest

import torch

from accelerate import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils import require_cpu





@require_cpu
class      lowerCAmelCase						(					unittest.TestCase			):



     def 			a__  (					self       ):
            _A=	torch.nn.Linear(10				,      10       )
            _A=	torch.optim.SGD(model.parameters()				,      0.1       )
            _A=	Accelerator()
            _A=	accelerator.prepare(lowerCAmelCase__       )
            try:
                   pickle.loads(pickle.dumps(lowerCAmelCase__       )       )
            except Exception as e:
                   self.fail(f"Accelerated optimizer pickling failed with {e}"       )
            AcceleratorState._reset_state()