





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.models.esm.modeling_esmfold import EsmForProteinFolding



class    _SCREAMING_SNAKE_CASE						:







 '''simple docstring'''

 def __init__(self				:   Optional[Any]       , UpperCAmelCase_				:   Dict       , UpperCAmelCase_				:   Tuple=13       , UpperCAmelCase_				:   int=7       , UpperCAmelCase_				:   Dict=False       , UpperCAmelCase_				:   Dict=True       , UpperCAmelCase_				:   Optional[int]=False       , UpperCAmelCase_				:   List[Any]=False       , UpperCAmelCase_				:   Union[str, Any]=19       , UpperCAmelCase_				:   Dict=32       , UpperCAmelCase_				:   Optional[Any]=5       , UpperCAmelCase_				:   Any=4       , UpperCAmelCase_				:   int=37       , UpperCAmelCase_				:   Union[str, Any]="gelu"       , UpperCAmelCase_				:   List[Any]=0.1       , UpperCAmelCase_				:   Tuple=0.1       , UpperCAmelCase_				:   Any=512       , UpperCAmelCase_				:   Tuple=16       , UpperCAmelCase_				:   Optional[Any]=2       , UpperCAmelCase_				:   Optional[Any]=0.02       , UpperCAmelCase_				:   str=3       , UpperCAmelCase_				:   Optional[int]=4       , UpperCAmelCase_				:   Tuple=None       , )       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =parent
   lowerCamelCase__:      List[str]			  =batch_size
   lowerCamelCase__:      Union[str, Any]			  =seq_length
   lowerCamelCase__:      List[str]			  =is_training
   lowerCamelCase__:      Optional[int]			  =use_input_mask
   lowerCamelCase__:      int			  =use_token_type_ids
   lowerCamelCase__:      int			  =use_labels
   lowerCamelCase__:      Any			  =vocab_size
   lowerCamelCase__:      Optional[Any]			  =hidden_size
   lowerCamelCase__:      Optional[Any]			  =num_hidden_layers
   lowerCamelCase__:      Union[str, Any]			  =num_attention_heads
   lowerCamelCase__:      Optional[Any]			  =intermediate_size
   lowerCamelCase__:      int			  =hidden_act
   lowerCamelCase__:      List[str]			  =hidden_dropout_prob
   lowerCamelCase__:      List[Any]			  =attention_probs_dropout_prob
   lowerCamelCase__:      Optional[Any]			  =max_position_embeddings
   lowerCamelCase__:      str			  =type_vocab_size
   lowerCamelCase__:      List[str]			  =type_sequence_label_size
   lowerCamelCase__:      Dict			  =initializer_range
   lowerCamelCase__:      Any			  =num_labels
   lowerCamelCase__:      int			  =num_choices
   lowerCamelCase__:      Tuple			  =scope

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->int:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =ids_tensor([self.batch_size, self.seq_length]       , self.vocab_size)

   lowerCamelCase__:      Optional[Any]			  =None
   if self.use_input_mask:
     lowerCamelCase__:      Tuple			  =random_attention_mask([self.batch_size, self.seq_length])

   lowerCamelCase__:      Optional[int]			  =None
   lowerCamelCase__:      Optional[int]			  =None
   lowerCamelCase__:      Any			  =None
   if self.use_labels:
     lowerCamelCase__:      Optional[Any]			  =ids_tensor([self.batch_size]       , self.type_sequence_label_size)
     lowerCamelCase__:      List[str]			  =ids_tensor([self.batch_size, self.seq_length]       , self.num_labels)
     lowerCamelCase__:      List[str]			  =ids_tensor([self.batch_size]       , self.num_choices)

   lowerCamelCase__:      Optional[Any]			  =self.get_config()

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

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =EsmConfig(
       vocab_size=33       , 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       , is_folding_model=UpperCAmelCase_       , esmfold_config={"trunk": {"num_blocks": 2}, "fp16_esm": False}       , )
   return config

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any]       , UpperCAmelCase_				:   int       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   str       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   str)       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      str			  =EsmForProteinFolding(config=UpperCAmelCase_).float()
   model.to(UpperCAmelCase_)
   model.eval()
   lowerCamelCase__:      str			  =model(UpperCAmelCase_       , attention_mask=UpperCAmelCase_)
   lowerCamelCase__:      Optional[int]			  =model(UpperCAmelCase_)
   lowerCamelCase__:      Any			  =model(UpperCAmelCase_)

   self.parent.assertEqual(result.positions.shape       , (8, self.batch_size, self.seq_length, 14, 3))
   self.parent.assertEqual(result.angles.shape       , (8, self.batch_size, self.seq_length, 7, 2))







 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->List[Any]:



   '''simple docstring'''



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



@require_torch
class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							unittest.TestCase    ):







 '''simple docstring'''
 lowercase_		=			False

 lowercase_		=			(EsmForProteinFolding,) if is_torch_available() else ()
 lowercase_		=			()
 lowercase_		=			{} if is_torch_available() else {}
 lowercase_		=			False

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =EsmFoldModelTester(self)
   lowerCamelCase__:      str			  =ConfigTester(self       , config_class=UpperCAmelCase_       , hidden_size=37)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->List[Any]:



   '''simple docstring'''



   self.config_tester.run_common_tests()

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*UpperCAmelCase_)

 @unittest.skip("Does not support attention outputs")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Any:



   '''simple docstring'''



   pass

 @unittest.skip
 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Tuple:



   '''simple docstring'''



   pass

 @unittest.skip("Esm does not support embedding resizing")
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->str:



   '''simple docstring'''



   pass

 @unittest.skip("Esm does not support embedding resizing")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->List[Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support passing input embeds!")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->Dict:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support head pruning.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Union[str, Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support head pruning.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->Tuple:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support head pruning.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Any:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support head pruning.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Any:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support head pruning.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->Optional[int]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not output hidden states in the normal way.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Optional[Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMfold does not output hidden states in the normal way.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->List[Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold only has one output format.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->Any:



   '''simple docstring'''



   pass

 @unittest.skip("This test doesn't work for ESMFold and doesn't test core functionality")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Any:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support input chunking.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->Optional[Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Union[str, Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold doesn't support torchscript compilation.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->List[str]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold doesn't support torchscript compilation.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->int:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold doesn't support torchscript compilation.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Optional[int]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold doesn't support data parallel.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->Union[str, Any]:



   '''simple docstring'''



   pass







 @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->int:



   '''simple docstring'''



   pass



@require_torch
class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ):







 '''simple docstring'''

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =EsmForProteinFolding.from_pretrained("facebook/esmfold_v1").float()
   model.eval()
   lowerCamelCase__:      Any			  =torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]])
   lowerCamelCase__:      int			  =model(UpperCAmelCase_)["positions"]
   lowerCamelCase__:      Tuple			  =torch.tensor([2.5828, 0.7993, -10.9334]       , dtype=torch.floataa)
   self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]       , UpperCAmelCase_       , atol=1E-4))








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

__A									=				logging.get_logger(__name__)



@add_end_docstrings(__SCREAMING_SNAKE_CASE    )
class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ):







 '''simple docstring'''

 def __init__(self				:   Tuple       , **UpperCAmelCase_				:   Tuple)       ->Any:



   '''simple docstring'''



   super().__init__(**UpperCAmelCase_)

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

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

 def __call__(self				:   Optional[int]       , UpperCAmelCase_				:   Union[str, "Image.Image", List[Dict[str, Any]]]       , UpperCAmelCase_				:   Union[str, List[str]] = None       , **UpperCAmelCase_				:   List[str]       , )       ->Union[str, Any]:



   '''simple docstring'''



   if "text_queries" in kwargs:
     lowerCamelCase__:      Any			  =kwargs.pop("text_queries")

   if isinstance(UpperCAmelCase_       , (str, Image.Image)):
     lowerCamelCase__:      List[Any]			  ={"image": image, "candidate_labels": candidate_labels}
   else:
     lowerCamelCase__:      Any			  =image
   lowerCamelCase__:      Dict			  =super().__call__(UpperCAmelCase_       , **UpperCAmelCase_)
   return results

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int]       , **UpperCAmelCase_				:   Union[str, Any])       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      List[str]			  ={}
   if "threshold" in kwargs:
     lowerCamelCase__:      List[Any]			  =kwargs["threshold"]
   if "top_k" in kwargs:
     lowerCamelCase__:      Any			  =kwargs["top_k"]
   return {}, {}, postprocess_params

 def 							SCREAMING_SNAKE_CASE_					(self				:   str       , UpperCAmelCase_				:   List[Any])       ->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:      List[Any]			  =load_image(inputs["image"])
   lowerCamelCase__:      Dict			  =inputs["candidate_labels"]
   if isinstance(UpperCAmelCase_       , UpperCAmelCase_):
     lowerCamelCase__:      Any			  =candidate_labels.split(",")

   lowerCamelCase__:      Optional[int]			  =torch.tensor([[image.height, image.width]]       , dtype=torch.intaa)
   for i, candidate_label in enumerate(UpperCAmelCase_):
     lowerCamelCase__:      Dict			  =self.tokenizer(UpperCAmelCase_       , return_tensors=self.framework)
     lowerCamelCase__:      Union[str, Any]			  =self.image_processor(UpperCAmelCase_       , return_tensors=self.framework)
     yield {
         "is_last": i == len(UpperCAmelCase_) - 1,
         "target_size": target_size,
         "candidate_label": candidate_label,
         **text_inputs,
         **image_features,
     }

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any]       , UpperCAmelCase_				:   Tuple)       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =model_inputs.pop("target_size")
   lowerCamelCase__:      Dict			  =model_inputs.pop("candidate_label")
   lowerCamelCase__:      Dict			  =model_inputs.pop("is_last")

   lowerCamelCase__:      Union[str, Any]			  =self.model(**UpperCAmelCase_)

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

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int]       , UpperCAmelCase_				:   int       , UpperCAmelCase_				:   Any=0.1       , UpperCAmelCase_				:   str=None)       ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:      Union[str, Any]			  =[]
   for model_output in model_outputs:
     lowerCamelCase__:      Optional[Any]			  =model_output["candidate_label"]
     lowerCamelCase__:      Tuple			  =BaseModelOutput(UpperCAmelCase_)
     lowerCamelCase__:      Dict			  =self.image_processor.post_process_object_detection(
         outputs=UpperCAmelCase_       , threshold=UpperCAmelCase_       , target_sizes=model_output["target_size"])[0]

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

       lowerCamelCase__:      Optional[Any]			  ={"score": score, "label": label, "box": box}
       results.append(UpperCAmelCase_)

   lowerCamelCase__:      List[str]			  =sorted(UpperCAmelCase_       , key=lambda UpperCAmelCase_: x["score"]       , reverse=UpperCAmelCase_)
   if top_k:
     lowerCamelCase__:      Dict			  =results[:top_k]

   return results







 def 							SCREAMING_SNAKE_CASE_					(self				:   str       , UpperCAmelCase_				:   "torch.Tensor")       ->Dict[str, int]:



   '''simple docstring'''



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

