



"""simple docstring"""
import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
       from transformers.models.esm.modeling_esm import (
           ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
           EsmEmbeddings,
           create_position_ids_from_input_ids,
       )


class    SCREAMING_SNAKE_CASE__ :

    def __init__(				self			:				Optional[int] ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				List[str]=1_3 ,  lowerCAmelCase_			:				Dict=7 ,  lowerCAmelCase_			:				Union[str, Any]=False ,  lowerCAmelCase_			:				Any=True ,  lowerCAmelCase_			:				Dict=False ,  lowerCAmelCase_			:				Optional[Any]=True ,  lowerCAmelCase_			:				Tuple=3_3 ,  lowerCAmelCase_			:				Dict=3_2 ,  lowerCAmelCase_			:				Optional[int]=5 ,  lowerCAmelCase_			:				Tuple=4 ,  lowerCAmelCase_			:				Tuple=3_7 ,  lowerCAmelCase_			:				Optional[Any]="gelu" ,  lowerCAmelCase_			:				Optional[Any]=0.1 ,  lowerCAmelCase_			:				Any=0.1 ,  lowerCAmelCase_			:				Tuple=5_1_2 ,  lowerCAmelCase_			:				int=1_6 ,  lowerCAmelCase_			:				Any=2 ,  lowerCAmelCase_			:				Tuple=0.02 ,  lowerCAmelCase_			:				Optional[Any]=3 ,  lowerCAmelCase_			:				int=4 ,  lowerCAmelCase_			:				Any=None ,  ):







      """simple docstring"""







      lowercase_								= parent
      lowercase_								= batch_size
      lowercase_								= seq_length
      lowercase_								= is_training
      lowercase_								= use_input_mask
      lowercase_								= use_token_type_ids
      lowercase_								= use_labels
      lowercase_								= vocab_size
      lowercase_								= hidden_size
      lowercase_								= num_hidden_layers
      lowercase_								= num_attention_heads
      lowercase_								= intermediate_size
      lowercase_								= hidden_act
      lowercase_								= hidden_dropout_prob
      lowercase_								= attention_probs_dropout_prob
      lowercase_								= max_position_embeddings
      lowercase_								= type_vocab_size
      lowercase_								= type_sequence_label_size
      lowercase_								= initializer_range
      lowercase_								= num_labels
      lowercase_								= num_choices
      lowercase_								= scope

    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size)

      lowercase_								= None
      if self.use_input_mask:
        lowercase_								= random_attention_mask([self.batch_size, self.seq_length])

      lowercase_								= None
      lowercase_								= None
      lowercase_								= None
      if self.use_labels:
        lowercase_								= ids_tensor([self.batch_size] ,  self.type_sequence_label_size)
        lowercase_								= ids_tensor([self.batch_size, self.seq_length] ,  self.num_labels)
        lowercase_								= ids_tensor([self.batch_size] ,  self.num_choices)

      lowercase_								= self.get_config()

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

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      return EsmConfig(
          vocab_size=self.vocab_size ,  hidden_size=self.hidden_size ,  pad_token_id=1 ,  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 ,  max_position_embeddings=self.max_position_embeddings ,  type_vocab_size=self.type_vocab_size ,  initializer_range=self.initializer_range ,  )

    def 					_UpperCAmelCase	(				self			:				Any ,  lowerCAmelCase_			:				Union[str, Any] ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				List[Any] ,  lowerCAmelCase_			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= EsmModel(config=lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_)
      lowercase_								= model(lowerCAmelCase_)
      lowercase_								= model(lowerCAmelCase_)

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

    def 					_UpperCAmelCase	(				self			:				Any ,  lowerCAmelCase_			:				Union[str, Any] ,  lowerCAmelCase_			:				List[str] ,  lowerCAmelCase_			:				List[Any] ,  lowerCAmelCase_			:				List[Any] ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Dict):







      """simple docstring"""







      lowercase_								= EsmForMaskedLM(config=lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  labels=lowerCAmelCase_)
      self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.seq_length, self.vocab_size))

    def 					_UpperCAmelCase	(				self			:				str ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				List[Any] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Tuple):







      """simple docstring"""







      lowercase_								= self.num_labels
      lowercase_								= EsmForTokenClassification(config=lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  labels=lowerCAmelCase_)
      self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.seq_length, self.num_labels))




    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= self.prepare_config_and_inputs()
      (
          (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    
      )								= config_and_inputs
      lowercase_								= {"""input_ids""": input_ids, """attention_mask""": input_mask}
      return config, inputs_dict







@require_torch
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase      ,  unittest.TestCase   ):
    lowercase__      =   False

    lowercase__      =   (
        (
            EsmForMaskedLM,
            EsmModel,
            EsmForSequenceClassification,
            EsmForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    lowercase__      =   ()
    lowercase__      =   (
        {
            "feature-extraction": EsmModel,
            "fill-mask": EsmForMaskedLM,
            "text-classification": EsmForSequenceClassification,
            "token-classification": EsmForTokenClassification,
            "zero-shot": EsmForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowercase__      =   True

    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= EsmModelTester(self)
      lowercase_								= ConfigTester(self ,  config_class=lowerCAmelCase_ ,  hidden_size=3_7)

    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      self.config_tester.run_common_tests()

    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				Optional[int]):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
        lowercase_								= type
        self.model_tester.create_and_check_model(*lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				Union[str, Any]):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase_)

    @slow
    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowercase_								= EsmModel.from_pretrained(lowerCAmelCase_)
        self.assertIsNotNone(lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				Optional[int]):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()[0]
      lowercase_								= EsmEmbeddings(config=lowerCAmelCase_)

      lowercase_								= torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]])
      lowercase_								= torch.as_tensor(
          [
              [
                  0 + model.padding_idx + 1,
                  1 + model.padding_idx + 1,
                  2 + model.padding_idx + 1,
                  model.padding_idx,
              ]
          ])
      lowercase_								= create_position_ids_from_input_ids(lowerCAmelCase_ ,  model.padding_idx)
      self.assertEqual(position_ids.shape ,  expected_positions.shape)
      self.assertTrue(torch.all(torch.eq(lowerCAmelCase_ ,  lowerCAmelCase_)))

    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()[0]
      lowercase_								= EsmEmbeddings(config=lowerCAmelCase_)

      lowercase_								= torch.empty(2 ,  4 ,  3_0)
      lowercase_								= [
          0 + embeddings.padding_idx + 1,
          1 + embeddings.padding_idx + 1,
          2 + embeddings.padding_idx + 1,
          3 + embeddings.padding_idx + 1,
      ]
      lowercase_								= torch.as_tensor([expected_single_positions, expected_single_positions])
      lowercase_								= embeddings.create_position_ids_from_inputs_embeds(lowerCAmelCase_)
      self.assertEqual(position_ids.shape ,  expected_positions.shape)
      self.assertTrue(torch.all(torch.eq(lowerCAmelCase_ ,  lowerCAmelCase_)))

    @unittest.skip("""Esm does not support embedding resizing""")
    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      pass

    @unittest.skip("""Esm does not support embedding resizing""")
    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      pass




    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests.""")
    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      pass







@require_torch
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase   ):

    @slow
    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      with torch.no_grad():
        lowercase_								= EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D""")
        model.eval()
        lowercase_								= torch.tensor([[0, 1, 2, 3, 4, 5]])
        lowercase_								= model(lowerCAmelCase_)[0]

        lowercase_								= 3_3

        lowercase_								= torch.Size((1, 6, vocab_size))
        self.assertEqual(output.shape ,  lowerCAmelCase_)

        lowercase_								= torch.tensor(
            [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]])
        self.assertTrue(torch.allclose(output[:, :3, :3] ,  lowerCAmelCase_ ,  atol=1E-4))




    @slow
    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      with torch.no_grad():
        lowercase_								= EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D""")
        model.eval()

        lowercase_								= torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]])
        lowercase_								= model(lowerCAmelCase_)[0]
        # compare the actual values for a slice.
        lowercase_								= torch.tensor(
            [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]])
        self.assertTrue(torch.allclose(output[:, :3, :3] ,  lowerCAmelCase_ ,  atol=1E-4))






"""simple docstring"""
from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
       from PIL import Image

       from ..image_utils import load_image

if is_torch_available():
       import torch

       from transformers.modeling_outputs import BaseModelOutput

       from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

UpperCAmelCase						:		Optional[Any]			= logging.get_logger(__name__)


@add_end_docstrings(__UpperCAmelCase   )
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase   ):

    def __init__(				self			:				Optional[int] ,  **lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)

      if self.framework == "tf":
        raise ValueError(F'''The {self.__class__} is only available in PyTorch.''')

      requires_backends(self ,  """vision""")
      self.check_model_type(lowerCAmelCase_)

    def __call__(				self			:				Dict ,  lowerCAmelCase_			:				Union[str, "Image.Image", List[Dict[str, Any]]] ,  lowerCAmelCase_			:				Union[str, List[str]] = None ,  **lowerCAmelCase_			:				Dict ,  ):







      """simple docstring"""







      if "text_queries" in kwargs:
        lowercase_								= kwargs.pop("""text_queries""")

      if isinstance(lowerCAmelCase_ ,  (str, Image.Image)):
        lowercase_								= {"""image""": image, """candidate_labels""": candidate_labels}
      else:
        lowercase_								= image
      lowercase_								= super().__call__(lowerCAmelCase_ ,  **lowerCAmelCase_)
      return results

    def 					_UpperCAmelCase	(				self			:				Optional[Any] ,  **lowerCAmelCase_			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= {}
      if "threshold" in kwargs:
        lowercase_								= kwargs["""threshold"""]
      if "top_k" in kwargs:
        lowercase_								= kwargs["""top_k"""]
      return {}, {}, postprocess_params

    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      lowercase_								= load_image(inputs["""image"""])
      lowercase_								= inputs["""candidate_labels"""]
      if isinstance(lowerCAmelCase_ ,  lowerCAmelCase_):
        lowercase_								= candidate_labels.split(""",""")

      lowercase_								= torch.tensor([[image.height, image.width]] ,  dtype=torch.intaa)
      for i, candidate_label in enumerate(lowerCAmelCase_):
        lowercase_								= self.tokenizer(lowerCAmelCase_ ,  return_tensors=self.framework)
        lowercase_								= self.image_processor(lowerCAmelCase_ ,  return_tensors=self.framework)
        yield {
            "is_last": i == len(lowerCAmelCase_) - 1,
            "target_size": target_size,
            "candidate_label": candidate_label,
            **text_inputs,
            **image_features,
        }

    def 					_UpperCAmelCase	(				self			:				List[Any] ,  lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      lowercase_								= model_inputs.pop("""target_size""")
      lowercase_								= model_inputs.pop("""candidate_label""")
      lowercase_								= model_inputs.pop("""is_last""")

      lowercase_								= self.model(**lowerCAmelCase_)

      lowercase_								= {"""target_size""": target_size, """candidate_label""": candidate_label, """is_last""": is_last, **outputs}
      return model_outputs

    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				Tuple=0.1 ,  lowerCAmelCase_			:				List[str]=None):







      """simple docstring"""







      lowercase_								= []
      for model_output in model_outputs:
        lowercase_								= model_output["""candidate_label"""]
        lowercase_								= BaseModelOutput(lowerCAmelCase_)
        lowercase_								= self.image_processor.post_process_object_detection(
            outputs=lowerCAmelCase_ ,  threshold=lowerCAmelCase_ ,  target_sizes=model_output["""target_size"""])[0]

        for index in outputs["scores"].nonzero():
          lowercase_								= outputs["""scores"""][index].item()
          lowercase_								= self._get_bounding_box(outputs["""boxes"""][index][0])

          lowercase_								= {"""score""": score, """label""": label, """box""": box}
          results.append(lowerCAmelCase_)

      lowercase_								= sorted(lowerCAmelCase_ ,  key=lambda lowerCAmelCase_: x["score"] ,  reverse=lowerCAmelCase_)
      if top_k:
        lowercase_								= results[:top_k]

      return results




    def 					_UpperCAmelCase	(				self			:				Union[str, Any] ,  lowerCAmelCase_			:				"torch.Tensor"):







      """simple docstring"""







      if self.framework != "pt":
        raise ValueError("""The ZeroShotObjectDetectionPipeline is only available in PyTorch.""")
      lowercase_					,    lowercase_					,    lowercase_					,    lowercase_								= box.int().tolist()
      lowercase_								= {
          """xmin""": xmin,
          """ymin""": ymin,
          """xmax""": xmax,
          """ymax""": ymax,
      }
      return bbox

