

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     __snake_case    :
       """simple docstring"""







       def __init__(		self      , _UpperCamelCase      , _UpperCamelCase=13      , _UpperCamelCase=7      , _UpperCamelCase=False      , _UpperCamelCase=True      , _UpperCamelCase=False      , _UpperCamelCase=True      , _UpperCamelCase=33      , _UpperCamelCase=32      , _UpperCamelCase=5      , _UpperCamelCase=4      , _UpperCamelCase=37      , _UpperCamelCase="gelu"      , _UpperCamelCase=0.1      , _UpperCamelCase=0.1      , _UpperCamelCase=5_12      , _UpperCamelCase=16      , _UpperCamelCase=2      , _UpperCamelCase=0.02      , _UpperCamelCase=3      , _UpperCamelCase=4      , _UpperCamelCase=None      , )      ->							str:






              """simple docstring"""


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


       def 			a		(		self   )      ->							Optional[int]:






              """simple docstring"""


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

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

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

              __snake_case       	=							self.get_config()

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


       def 			a		(		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 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							Optional[int]:






              """simple docstring"""


              __snake_case       	=							EsmModel(config=_UpperCamelCase   )
              model.to(_UpperCamelCase   )
              model.eval()
              __snake_case       	=							model(_UpperCamelCase      , attention_mask=_UpperCamelCase   )
              __snake_case       	=							model(_UpperCamelCase   )
              __snake_case       	=							model(_UpperCamelCase   )

              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 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							Optional[Any]:






              """simple docstring"""


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


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							Dict:






              """simple docstring"""


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






       def 			a		(		self   )      ->							Optional[int]:






              """simple docstring"""


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






@require_torch
class     __snake_case    (				snake_case__			,	snake_case__			,	unittest.TestCase  ):
       """simple docstring"""







       __SCREAMING_SNAKE_CASE				    =							False

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


       def 			a		(		self   )      ->							Any:






              """simple docstring"""


              __snake_case       	=							EsmModelTester(self   )
              __snake_case       	=							ConfigTester(self      , config_class=_UpperCamelCase      , hidden_size=37   )


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






              """simple docstring"""


              self.config_tester.run_common_tests()


       def 			a		(		self   )      ->							Optional[int]:






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_model(*_UpperCamelCase   )


       def 			a		(		self   )      ->							Tuple:






              """simple docstring"""


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


       def 			a		(		self   )      ->							Tuple:






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase   )


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






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase   )


       @slow
       def 			a		(		self   )      ->							Optional[Any]:






              """simple docstring"""


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


       def 			a		(		self   )      ->							Tuple:






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()[0]
              __snake_case       	=							EsmEmbeddings(config=_UpperCamelCase   )

              __snake_case       	=							torch.as_tensor([[12, 31, 13, model.padding_idx]]   )
              __snake_case       	=							torch.as_tensor(
                  [
                      [
                          0 + model.padding_idx + 1,
                          1 + model.padding_idx + 1,
                          2 + model.padding_idx + 1,
                          model.padding_idx,
                      ]
                  ]   )
              __snake_case       	=							create_position_ids_from_input_ids(_UpperCamelCase      , model.padding_idx   )
              self.assertEqual(position_ids.shape      , expected_positions.shape   )
              self.assertTrue(torch.all(torch.eq(_UpperCamelCase      , _UpperCamelCase   )   )   )


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






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()[0]
              __snake_case       	=							EsmEmbeddings(config=_UpperCamelCase   )

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


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






              """simple docstring"""


              pass


       @unittest.skip("""Esm does not support embedding resizing"""   )
       def 			a		(		self   )      ->							Union[str, Any]:






              """simple docstring"""


              pass






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






              """simple docstring"""


              pass






@require_torch
class     __snake_case    (				snake_case__  ):
       """simple docstring"""







       @slow
       def 			a		(		self   )      ->							Union[str, Any]:






              """simple docstring"""


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

                     __snake_case       	=							33

                     __snake_case       	=							torch.Size((1, 6, vocab_size)   )
                     self.assertEqual(output.shape      , _UpperCamelCase   )

                     __snake_case       	=							torch.tensor(
                         [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]   )
                     self.assertTrue(torch.allclose(output[:, :3, :3]      , _UpperCamelCase      , atol=1E-4   )   )






       @slow
       def 			a		(		self   )      ->							Optional[int]:






              """simple docstring"""


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

                     __snake_case       	=							torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]   )
                     __snake_case       	=							model(_UpperCamelCase   )[0]
                     # compare the actual values for a slice.
                     __snake_case       	=							torch.tensor(
                         [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]   )
                     self.assertTrue(torch.allclose(output[:, :3, :3]      , _UpperCamelCase      , atol=1E-4   )   )




import sys
from typing import Tuple

import numpy as np
import torch
from PIL import Image
from torch import nn

from transformers.image_utils import PILImageResampling
from utils import img_tensorize

class     __snake_case    :
       """simple docstring"""







       def __init__(		self      , _UpperCamelCase      , _UpperCamelCase=sys.maxsize   )      ->							Any:






              """simple docstring"""


              __snake_case       	=							"""bilinear"""
              __snake_case       	=							max_size
              __snake_case       	=							short_edge_length






       def __call__(		self      , _UpperCamelCase   )      ->							Optional[Any]:






              """simple docstring"""


              __snake_case       	=							[]
              for img in imgs:
                     __snake_case	,    __snake_case       	=							img.shape[:2]
                     # later: provide list and randomly choose index for resize
                     __snake_case       	=							np.random.randint(self.short_edge_length[0]      , self.short_edge_length[1] + 1   )
                     if size == 0:
                            return img
                     __snake_case       	=							size * 1.0 / min(_UpperCamelCase      , _UpperCamelCase   )
                     if h < w:
                            __snake_case	,    __snake_case       	=							size, scale * w
                     else:
                            __snake_case	,    __snake_case       	=							scale * h, size
                     if max(_UpperCamelCase      , _UpperCamelCase   ) > self.max_size:
                            __snake_case       	=							self.max_size * 1.0 / max(_UpperCamelCase      , _UpperCamelCase   )
                            __snake_case       	=							newh * scale
                            __snake_case       	=							neww * scale
                     __snake_case       	=							int(neww + 0.5   )
                     __snake_case       	=							int(newh + 0.5   )

                     if img.dtype == np.uinta:
                            __snake_case       	=							Image.fromarray(_UpperCamelCase   )
                            __snake_case       	=							pil_image.resize((neww, newh)      , PILImageResampling.BILINEAR   )
                            __snake_case       	=							np.asarray(_UpperCamelCase   )
                     else:
                            __snake_case       	=							img.permute(2      , 0      , 1   ).unsqueeze(0   )  # 3, 0, 1)  # hw(c) -> nchw
                            __snake_case       	=							nn.functional.interpolate(
                                _UpperCamelCase      , (newh, neww)      , mode=self.interp_method      , align_corners=_UpperCamelCase   ).squeeze(0   )
                     img_augs.append(_UpperCamelCase   )

              return img_augs






class     __snake_case    :
       """simple docstring"""







       def __init__(		self      , _UpperCamelCase   )      ->							Union[str, Any]:






              """simple docstring"""


              __snake_case       	=							ResizeShortestEdge([cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST]      , cfg.INPUT.MAX_SIZE_TEST   )
              __snake_case       	=							cfg.INPUT.FORMAT
              __snake_case       	=							cfg.SIZE_DIVISIBILITY
              __snake_case       	=							cfg.PAD_VALUE
              __snake_case       	=							cfg.INPUT.MAX_SIZE_TEST
              __snake_case       	=							cfg.MODEL.DEVICE
              __snake_case       	=							torch.tensor(cfg.MODEL.PIXEL_STD   ).to(self.device   ).view(len(cfg.MODEL.PIXEL_STD   )      , 1      , 1   )
              __snake_case       	=							torch.tensor(cfg.MODEL.PIXEL_MEAN   ).to(self.device   ).view(len(cfg.MODEL.PIXEL_STD   )      , 1      , 1   )
              __snake_case       	=							lambda _UpperCamelCase   : (x - self.pixel_mean) / self.pixel_std


       def 			a		(		self      , _UpperCamelCase   )      ->							int:






              """simple docstring"""


              __snake_case       	=							tuple(max(_UpperCamelCase   ) for s in zip(*[img.shape for img in images]   )   )
              __snake_case       	=							[im.shape[-2:] for im in images]
              __snake_case       	=							[
                  nn.functional.pad(
                      _UpperCamelCase      , [0, max_size[-1] - size[1], 0, max_size[-2] - size[0]]      , value=self.pad_value      , )
                  for size, im in zip(_UpperCamelCase      , _UpperCamelCase   )
              ]

              return torch.stack(_UpperCamelCase   ), torch.tensor(_UpperCamelCase   )






       def __call__(		self      , _UpperCamelCase      , _UpperCamelCase=False   )      ->							Union[str, Any]:






              """simple docstring"""


              with torch.no_grad():
                     if not isinstance(_UpperCamelCase      , _UpperCamelCase   ):
                            __snake_case       	=							[images]
                     if single_image:
                            assert len(_UpperCamelCase   ) == 1
                     for i in range(len(_UpperCamelCase   )   ):
                            if isinstance(images[i]      , torch.Tensor   ):
                                   images.insert(_UpperCamelCase      , images.pop(_UpperCamelCase   ).to(self.device   ).float()   )
                            elif not isinstance(images[i]      , torch.Tensor   ):
                                   images.insert(
                                       _UpperCamelCase      , torch.as_tensor(img_tensorize(images.pop(_UpperCamelCase   )      , input_format=self.input_format   )   )
                                       .to(self.device   )
                                       .float()      , )
            # resize smallest edge
                     __snake_case       	=							torch.tensor([im.shape[:2] for im in images]   )
                     __snake_case       	=							self.aug(_UpperCamelCase   )
                     # transpose images and convert to torch tensors
                     # images = [torch.as_tensor(i.astype("float32")).permute(2, 0, 1).to(self.device) for i in images]
                     # now normalize before pad to avoid useless arithmetic
                     __snake_case       	=							[self.normalizer(_UpperCamelCase   ) for x in images]
                     # now pad them to do the following operations
                     __snake_case	,    __snake_case       	=							self.pad(_UpperCamelCase   )
                     # Normalize

                     if self.size_divisibility > 0:
                            raise NotImplementedError()
                     # pad
                     __snake_case       	=							torch.true_divide(_UpperCamelCase      , _UpperCamelCase   )
                     if single_image:
                            return images[0], sizes[0], scales_yx[0]
                     else:
                            return images, sizes, scales_yx







def 						lowerCamelCase__      ( __A    :Union[str, Any]	,__A    :List[Any]		):


       """simple docstring"""




       boxes[:, 0::2] *= scale_yx[:, 1]
       boxes[:, 1::2] *= scale_yx[:, 0]
       return boxes







def 						lowerCamelCase__      ( __A    :Tuple	,__A    :Tuple[int, int]		):


       """simple docstring"""




       assert torch.isfinite(__A		).all(), "Box tensor contains infinite or NaN!"
       __snake_case	,    __snake_case       	=							box_size
       tensor[:, 0].clamp_(min=0	,max=__A		)
       tensor[:, 1].clamp_(min=0	,max=__A		)
       tensor[:, 2].clamp_(min=0	,max=__A		)
       tensor[:, 3].clamp_(min=0	,max=__A		)

