






from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers.models.distilbert.modeling_tf_distilbert import (
     TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
     TFDistilBertForMaskedLM,
     TFDistilBertForMultipleChoice,
     TFDistilBertForQuestionAnswering,
     TFDistilBertForSequenceClassification,
     TFDistilBertForTokenClassification,
     TFDistilBertModel,
 )





class       __a  :







 def __init__(		self :       Optional[Any]							,       SCREAMING_SNAKE_CASE :       Tuple							,       ):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  parent
      UpperCamelCase__      :       Union[str, Any]	    =  13
      UpperCamelCase__      :       Optional[Any]	    =  7
      UpperCamelCase__      :       str	    =  True
      UpperCamelCase__      :       List[Any]	    =  True
      UpperCamelCase__      :       Dict	    =  False
      UpperCamelCase__      :       int	    =  True
      UpperCamelCase__      :       Tuple	    =  99
      UpperCamelCase__      :       Tuple	    =  32
      UpperCamelCase__      :       int	    =  2
      UpperCamelCase__      :       List[str]	    =  4
      UpperCamelCase__      :       Union[str, Any]	    =  37
      UpperCamelCase__      :       Tuple	    =  "gelu"
      UpperCamelCase__      :       Any	    =  0.1
      UpperCamelCase__      :       Optional[int]	    =  0.1
      UpperCamelCase__      :       int	    =  5_12
      UpperCamelCase__      :       Dict	    =  16
      UpperCamelCase__      :       Union[str, Any]	    =  2
      UpperCamelCase__      :       List[str]	    =  0.0_2
      UpperCamelCase__      :       str	    =  3
      UpperCamelCase__      :       Tuple	    =  4
      UpperCamelCase__      :       Union[str, Any]	    =  None







 def  __lowercase (		self :       Any						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[int]	    =  ids_tensor([self.batch_size, self.seq_length]							,       self.vocab_size						)

      UpperCamelCase__      :       Optional[int]	    =  None
      if self.use_input_mask:
           UpperCamelCase__      :       int	    =  random_attention_mask([self.batch_size, self.seq_length]						)

      UpperCamelCase__      :       List[Any]	    =  None
      UpperCamelCase__      :       str	    =  None
      UpperCamelCase__      :       Tuple	    =  None
      if self.use_labels:
           UpperCamelCase__      :       str	    =  ids_tensor([self.batch_size]							,       self.type_sequence_label_size						)
           UpperCamelCase__      :       Union[str, Any]	    =  ids_tensor([self.batch_size, self.seq_length]							,       self.num_labels						)
           UpperCamelCase__      :       Tuple	    =  ids_tensor([self.batch_size]							,       self.num_choices						)

      UpperCamelCase__      :       List[Any]	    =  DistilBertConfig(
          vocab_size=self.vocab_size							,       dim=self.hidden_size							,       n_layers=self.num_hidden_layers							,       n_heads=self.num_attention_heads							,       hidden_dim=self.intermediate_size							,       hidden_act=self.hidden_act							,       dropout=self.hidden_dropout_prob							,       attention_dropout=self.attention_probs_dropout_prob							,       max_position_embeddings=self.max_position_embeddings							,       initializer_range=self.initializer_range							,       )

      return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







 def  __lowercase (		self :       Optional[int]							,       SCREAMING_SNAKE_CASE :       Optional[Any]							,       SCREAMING_SNAKE_CASE :       Dict							,       SCREAMING_SNAKE_CASE :       int							,       SCREAMING_SNAKE_CASE :       List[str]							,       SCREAMING_SNAKE_CASE :       List[str]							,       SCREAMING_SNAKE_CASE :       int						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[Any]	    =  TFDistilBertModel(config=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__      :       Any	    =  {"input_ids": input_ids, "attention_mask": input_mask}

      UpperCamelCase__      :       Tuple	    =  model(SCREAMING_SNAKE_CASE						)

      UpperCamelCase__      :       Tuple	    =  [input_ids, input_mask]

      UpperCamelCase__      :       Optional[int]	    =  model(SCREAMING_SNAKE_CASE						)

      self.parent.assertEqual(result.last_hidden_state.shape							,       (self.batch_size, self.seq_length, self.hidden_size)						)







 def  __lowercase (		self :       Dict							,       SCREAMING_SNAKE_CASE :       Tuple							,       SCREAMING_SNAKE_CASE :       List[Any]							,       SCREAMING_SNAKE_CASE :       Any							,       SCREAMING_SNAKE_CASE :       int							,       SCREAMING_SNAKE_CASE :       str							,       SCREAMING_SNAKE_CASE :       Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       List[str]	    =  TFDistilBertForMaskedLM(config=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__      :       Optional[Any]	    =  {"input_ids": input_ids, "attention_mask": input_mask}
      UpperCamelCase__      :       str	    =  model(SCREAMING_SNAKE_CASE						)
      self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.seq_length, self.vocab_size)						)







 def  __lowercase (		self :       int							,       SCREAMING_SNAKE_CASE :       int							,       SCREAMING_SNAKE_CASE :       str							,       SCREAMING_SNAKE_CASE :       Optional[Any]							,       SCREAMING_SNAKE_CASE :       Optional[Any]							,       SCREAMING_SNAKE_CASE :       Dict							,       SCREAMING_SNAKE_CASE :       int						):


      '''simple docstring'''






      UpperCamelCase__      :       Any	    =  TFDistilBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__      :       Tuple	    =  {
          "input_ids": input_ids,
          "attention_mask": input_mask,
      }
      UpperCamelCase__      :       List[Any]	    =  model(SCREAMING_SNAKE_CASE						)
      self.parent.assertEqual(result.start_logits.shape							,       (self.batch_size, self.seq_length)						)
      self.parent.assertEqual(result.end_logits.shape							,       (self.batch_size, self.seq_length)						)







 def  __lowercase (		self :       Union[str, Any]							,       SCREAMING_SNAKE_CASE :       Any							,       SCREAMING_SNAKE_CASE :       Tuple							,       SCREAMING_SNAKE_CASE :       str							,       SCREAMING_SNAKE_CASE :       List[Any]							,       SCREAMING_SNAKE_CASE :       Tuple							,       SCREAMING_SNAKE_CASE :       List[Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       Any	    =  self.num_labels
      UpperCamelCase__      :       List[Any]	    =  TFDistilBertForSequenceClassification(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__      :       List[Any]	    =  {"input_ids": input_ids, "attention_mask": input_mask}
      UpperCamelCase__      :       str	    =  model(SCREAMING_SNAKE_CASE						)
      self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.num_labels)						)







 def  __lowercase (		self :       str							,       SCREAMING_SNAKE_CASE :       List[str]							,       SCREAMING_SNAKE_CASE :       List[str]							,       SCREAMING_SNAKE_CASE :       List[Any]							,       SCREAMING_SNAKE_CASE :       Union[str, Any]							,       SCREAMING_SNAKE_CASE :       Optional[int]							,       SCREAMING_SNAKE_CASE :       List[str]						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[int]	    =  self.num_choices
      UpperCamelCase__      :       Union[str, Any]	    =  TFDistilBertForMultipleChoice(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__      :       int	    =  tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE							,       1						)							,       (1, self.num_choices, 1)						)
      UpperCamelCase__      :       List[Any]	    =  tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE							,       1						)							,       (1, self.num_choices, 1)						)
      UpperCamelCase__      :       str	    =  {
          "input_ids": multiple_choice_inputs_ids,
          "attention_mask": multiple_choice_input_mask,
      }
      UpperCamelCase__      :       int	    =  model(SCREAMING_SNAKE_CASE						)
      self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.num_choices)						)







 def  __lowercase (		self :       int							,       SCREAMING_SNAKE_CASE :       Optional[int]							,       SCREAMING_SNAKE_CASE :       Tuple							,       SCREAMING_SNAKE_CASE :       Union[str, Any]							,       SCREAMING_SNAKE_CASE :       Tuple							,       SCREAMING_SNAKE_CASE :       int							,       SCREAMING_SNAKE_CASE :       Tuple						):


      '''simple docstring'''






      UpperCamelCase__      :       Any	    =  self.num_labels
      UpperCamelCase__      :       str	    =  TFDistilBertForTokenClassification(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__      :       int	    =  {"input_ids": input_ids, "attention_mask": input_mask}
      UpperCamelCase__      :       Union[str, Any]	    =  model(SCREAMING_SNAKE_CASE						)
      self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.seq_length, self.num_labels)						)







 def  __lowercase (		self :       List[str]						):


      '''simple docstring'''






      UpperCamelCase__      :       str	    =  self.prepare_config_and_inputs()
      ((UpperCamelCase__)    ,					(UpperCamelCase__)    ,					(UpperCamelCase__)    ,					(UpperCamelCase__)    ,					(UpperCamelCase__)    ,					(UpperCamelCase__))      :       Optional[int]	    =  config_and_inputs
      UpperCamelCase__      :       Union[str, Any]	    =  {"input_ids": input_ids, "attention_mask": input_mask}
      return config, inputs_dict







@require_tf
class       __a  (   A__	,       A__	,       unittest.TestCase	):
 _lowerCAmelCase	:	Optional[Any]				=			(
     (
         TFDistilBertModel,
         TFDistilBertForMaskedLM,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertForMultipleChoice,
     )
     if is_tf_available()
     else None
 )
 _lowerCAmelCase	:	Optional[Any]				=			(
     {
         '''feature-extraction''': TFDistilBertModel,
         '''fill-mask''': TFDistilBertForMaskedLM,
         '''question-answering''': TFDistilBertForQuestionAnswering,
         '''text-classification''': TFDistilBertForSequenceClassification,
         '''token-classification''': TFDistilBertForTokenClassification,
         '''zero-shot''': TFDistilBertForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 _lowerCAmelCase	:	List[str]				=			False
 _lowerCAmelCase	:	Optional[int]				=			False







 def  __lowercase (		self :       List[str]						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[int]	    =  TFDistilBertModelTester(self						)
      UpperCamelCase__      :       List[str]	    =  ConfigTester(self							,       config_class=SCREAMING_SNAKE_CASE							,       dim=37						)







 def  __lowercase (		self :       int						):


      '''simple docstring'''






      self.config_tester.run_common_tests()







 def  __lowercase (		self :       List[Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       Any	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_model(*SCREAMING_SNAKE_CASE						)







 def  __lowercase (		self :       Dict						):


      '''simple docstring'''






      UpperCamelCase__      :       List[Any]	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_masked_lm(*SCREAMING_SNAKE_CASE						)







 def  __lowercase (		self :       str						):


      '''simple docstring'''






      UpperCamelCase__      :       Union[str, Any]	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_question_answering(*SCREAMING_SNAKE_CASE						)







 def  __lowercase (		self :       Any						):


      '''simple docstring'''






      UpperCamelCase__      :       List[Any]	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_sequence_classification(*SCREAMING_SNAKE_CASE						)







 def  __lowercase (		self :       str						):


      '''simple docstring'''






      UpperCamelCase__      :       Dict	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_multiple_choice(*SCREAMING_SNAKE_CASE						)







 def  __lowercase (		self :       Any						):


      '''simple docstring'''






      UpperCamelCase__      :       str	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_distilbert_for_token_classification(*SCREAMING_SNAKE_CASE						)







 @slow
 def  __lowercase (		self :       Dict						):


      '''simple docstring'''






      for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]						):
           UpperCamelCase__      :       Any	    =  TFDistilBertModel.from_pretrained(SCREAMING_SNAKE_CASE						)
           self.assertIsNotNone(SCREAMING_SNAKE_CASE						)







@require_tf
class       __a  (   unittest.TestCase	):







 @slow
 def  __lowercase (		self :       Tuple						):


      '''simple docstring'''






      UpperCamelCase__      :       List[str]	    =  TFDistilBertModel.from_pretrained("distilbert-base-uncased"						)
      UpperCamelCase__      :       str	    =  tf.constant([[0, 1, 2, 3, 4, 5]]						)
      UpperCamelCase__      :       int	    =  model(SCREAMING_SNAKE_CASE						)[0]

      UpperCamelCase__      :       Any	    =  [1, 6, 7_68]
      self.assertEqual(output.shape							,       SCREAMING_SNAKE_CASE						)

      UpperCamelCase__      :       Union[str, Any]	    =  tf.constant(
          [
              [
                  [0.1_9_2_6_1_8_8_5, -0.1_3_7_3_2_9_5_5, 0.4_1_1_9_7_9_9],
                  [0.2_2_1_5_0_1_5_6, -0.0_7_4_2_2_6_6_1, 0.3_9_0_3_7_2_0_4],
                  [0.2_2_7_5_6_0_1_8, -0.0_8_9_6_4_1_4, 0.3_7_0_1_4_6_7],
              ]
          ]						)
      tf.debugging.assert_near(output[:, :3, :3]							,       SCREAMING_SNAKE_CASE							,       atol=1e-4						)







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

lowerCamelCase     :							str    =logging.get_logger(__name__)





@add_end_docstrings(A__	)
class       __a  (   A__	):







 def __init__(		self :       List[str]							,       **SCREAMING_SNAKE_CASE :       Tuple						):


      '''simple docstring'''






      super().__init__(**SCREAMING_SNAKE_CASE						)

      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 :       Dict							,       SCREAMING_SNAKE_CASE :       Union[str, List[str], "Image", List["Image"]]							,       **SCREAMING_SNAKE_CASE :       str						):


      '''simple docstring'''






      return super().__call__(SCREAMING_SNAKE_CASE							,       **SCREAMING_SNAKE_CASE						)







 def  __lowercase (		self :       List[str]							,       **SCREAMING_SNAKE_CASE :       List[str]						):


      '''simple docstring'''






      UpperCamelCase__      :       List[Any]	    =  {}
      if "candidate_labels" in kwargs:
           UpperCamelCase__      :       Optional[Any]	    =  kwargs["candidate_labels"]
      if "hypothesis_template" in kwargs:
           UpperCamelCase__      :       int	    =  kwargs["hypothesis_template"]

      return preprocess_params, {}, {}







 def  __lowercase (		self :       List[Any]							,       SCREAMING_SNAKE_CASE :       List[Any]							,       SCREAMING_SNAKE_CASE :       Tuple=None							,       SCREAMING_SNAKE_CASE :       Optional[int]="This is a photo of {}."						):


      '''simple docstring'''






      UpperCamelCase__      :       Dict	    =  load_image(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__      :       List[Any]	    =  self.image_processor(images=[image]							,       return_tensors=self.framework						)
      UpperCamelCase__      :       Any	    =  candidate_labels
      UpperCamelCase__      :       Dict	    =  [hypothesis_template.format(SCREAMING_SNAKE_CASE						) for x in candidate_labels]
      UpperCamelCase__      :       Optional[Any]	    =  self.tokenizer(SCREAMING_SNAKE_CASE							,       return_tensors=self.framework							,       padding=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__      :       List[Any]	    =  [text_inputs]
      return inputs







 def  __lowercase (		self :       int							,       SCREAMING_SNAKE_CASE :       Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       Tuple	    =  model_inputs.pop("candidate_labels"						)
      UpperCamelCase__      :       List[str]	    =  model_inputs.pop("text_inputs"						)
      if isinstance(text_inputs[0]							,       SCREAMING_SNAKE_CASE						):
           UpperCamelCase__      :       Dict	    =  text_inputs[0]
      else:
           # Batching case.
           UpperCamelCase__      :       Union[str, Any]	    =  text_inputs[0][0]

      UpperCamelCase__      :       Any	    =  self.model(**SCREAMING_SNAKE_CASE							,       **SCREAMING_SNAKE_CASE						)

      UpperCamelCase__      :       Dict	    =  {
          "candidate_labels": candidate_labels,
          "logits": outputs.logits_per_image,
      }
      return model_outputs







 def  __lowercase (		self :       Any							,       SCREAMING_SNAKE_CASE :       List[str]						):


      '''simple docstring'''






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

      UpperCamelCase__      :       Optional[int]	    =  [
          {"score": score, "label": candidate_label}
          for score, candidate_label in sorted(zip(SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE						)							,       key=lambda SCREAMING_SNAKE_CASE						: -x[0]						)
      ]
      return result