'''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 			UpperCAmelCase  :





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase=13	,   __lowerCAmelCase=7	,   __lowerCAmelCase=False	,   __lowerCAmelCase=True	,   __lowerCAmelCase=False	,   __lowerCAmelCase=True	,   __lowerCAmelCase=33	,   __lowerCAmelCase=32	,   __lowerCAmelCase=5	,   __lowerCAmelCase=4	,   __lowerCAmelCase=37	,   __lowerCAmelCase="gelu"	,   __lowerCAmelCase=0.1	,   __lowerCAmelCase=0.1	,   __lowerCAmelCase=512	,   __lowerCAmelCase=16	,   __lowerCAmelCase=2	,   __lowerCAmelCase=0.0_2	,   __lowerCAmelCase=3	,   __lowerCAmelCase=4	,   __lowerCAmelCase=None	,   )  ->  Optional[Any]:
        lowercase__				:    Dict  =					parent
        lowercase__				:    List[Any]  =					batch_size
        lowercase__				:    str  =					seq_length
        lowercase__				:    Optional[int]  =					is_training
        lowercase__				:    List[str]  =					use_input_mask
        lowercase__				:    Any  =					use_token_type_ids
        lowercase__				:    List[Any]  =					use_labels
        lowercase__				:    Union[str, Any]  =					vocab_size
        lowercase__				:    str  =					hidden_size
        lowercase__				:    Dict  =					num_hidden_layers
        lowercase__				:    str  =					num_attention_heads
        lowercase__				:    Optional[int]  =					intermediate_size
        lowercase__				:    Optional[Any]  =					hidden_act
        lowercase__				:    List[str]  =					hidden_dropout_prob
        lowercase__				:    Any  =					attention_probs_dropout_prob
        lowercase__				:    int  =					max_position_embeddings
        lowercase__				:    Optional[Any]  =					type_vocab_size
        lowercase__				:    Union[str, Any]  =					type_sequence_label_size
        lowercase__				:    str  =					initializer_range
        lowercase__				:    Union[str, Any]  =					num_labels
        lowercase__				:    str  =					num_choices
        lowercase__				:    Any  =					scope

       def _lowerCAmelCase(						self )  ->  Dict:
        lowercase__				:    Optional[Any]  =					ids_tensor([self.batch_size, self.seq_length]	,   self.vocab_size )

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

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

        lowercase__				:    Optional[int]  =					self.get_config()

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

       def _lowerCAmelCase(						self )  ->  List[str]:
        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 _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase )  ->  Any:
        lowercase__				:    List[Any]  =					EsmModel(config=__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()
        lowercase__				:    str  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase )
        lowercase__				:    Dict  =					model(__lowerCAmelCase )
        lowercase__				:    List[str]  =					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 _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase )  ->  Union[str, Any]:
        lowercase__				:    Any  =					EsmForMaskedLM(config=__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()
        lowercase__				:    Tuple  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase	,   labels=__lowerCAmelCase )
        self.parent.assertEqual(result.logits.shape	,   (self.batch_size, self.seq_length, self.vocab_size) )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase )  ->  List[str]:
        lowercase__				:    Union[str, Any]  =					self.num_labels
        lowercase__				:    str  =					EsmForTokenClassification(config=__lowerCAmelCase )
        model.to(__lowerCAmelCase )
        model.eval()
        lowercase__				:    Dict  =					model(__lowerCAmelCase	,   attention_mask=__lowerCAmelCase	,   labels=__lowerCAmelCase )
        self.parent.assertEqual(result.logits.shape	,   (self.batch_size, self.seq_length, self.num_labels) )


       def _lowerCAmelCase(						self )  ->  Dict:
        lowercase__				:    int  =					self.prepare_config_and_inputs()
        (
            (
            lowercase__
        ) ,							(
            lowercase__
        ) ,							(
            lowercase__
        ) ,							(
            lowercase__
        ) ,							(
            lowercase__
        ) ,							(
            lowercase__
        ) ,							
        )				:    List[str]  =					config_and_inputs
        lowercase__				:    Optional[Any]  =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
        return config, inputs_dict



@require_torch
class 			UpperCAmelCase  (						a__    ,						a__    ,						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 _lowerCAmelCase(						self )  ->  Tuple:
        lowercase__				:    Tuple  =					EsmModelTester(self )
        lowercase__				:    str  =					ConfigTester(self	,   config_class=__lowerCAmelCase	,   hidden_size=37 )

       def _lowerCAmelCase(						self )  ->  Dict:
        self.config_tester.run_common_tests()

       def _lowerCAmelCase(						self )  ->  Tuple:
        lowercase__				:    Tuple  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  Tuple:
        lowercase__				:    List[str]  =					self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
         lowercase__				:    Union[str, Any]  =					type
         self.model_tester.create_and_check_model(*__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  Tuple:
        lowercase__				:    Union[str, Any]  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  Union[str, Any]:
        lowercase__				:    Optional[Any]  =					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase )

       @slow
       def _lowerCAmelCase(						self )  ->  Union[str, Any]:
        for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowercase__				:    Any  =					EsmModel.from_pretrained(__lowerCAmelCase )
         self.assertIsNotNone(__lowerCAmelCase )

       def _lowerCAmelCase(						self )  ->  Dict:
        lowercase__				:    int  =					self.model_tester.prepare_config_and_inputs()[0]
        lowercase__				:    List[str]  =					EsmEmbeddings(config=__lowerCAmelCase )

        lowercase__				:    List[Any]  =					torch.as_tensor([[12, 31, 13, model.padding_idx]] )
        lowercase__				:    Union[str, Any]  =					torch.as_tensor(
            [
                [
                    0 + model.padding_idx + 1,
                    1 + model.padding_idx + 1,
                    2 + model.padding_idx + 1,
                    model.padding_idx,
                ]
            ] )
        lowercase__				:    List[Any]  =					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 _lowerCAmelCase(						self )  ->  Tuple:
        lowercase__				:    List[Any]  =					self.model_tester.prepare_config_and_inputs()[0]
        lowercase__				:    Optional[Any]  =					EsmEmbeddings(config=__lowerCAmelCase )

        lowercase__				:    str  =					torch.empty(2	,   4	,   30 )
        lowercase__				:    Optional[Any]  =					[
            0 + embeddings.padding_idx + 1,
            1 + embeddings.padding_idx + 1,
            2 + embeddings.padding_idx + 1,
            3 + embeddings.padding_idx + 1,
        ]
        lowercase__				:    Dict  =					torch.as_tensor([expected_single_positions, expected_single_positions] )
        lowercase__				:    int  =					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 _lowerCAmelCase(						self )  ->  Optional[int]:
        pass

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


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



@require_torch
class 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       @slow
       def _lowerCAmelCase(						self )  ->  List[str]:
        with torch.no_grad():
         lowercase__				:    List[Any]  =					EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D''' )
         model.eval()
         lowercase__				:    Any  =					torch.tensor([[0, 1, 2, 3, 4, 5]] )
         lowercase__				:    int  =					model(__lowerCAmelCase )[0]

         lowercase__				:    Union[str, Any]  =					33

         lowercase__				:    List[str]  =					torch.Size((1, 6, vocab_size) )
         self.assertEqual(output.shape	,   __lowerCAmelCase )

         lowercase__				:    List[str]  =					torch.tensor(
             [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]] )
         self.assertTrue(torch.allclose(output[:, :3, :3]	,   __lowerCAmelCase	,   atol=1E-4 ) )


       @slow
       def _lowerCAmelCase(						self )  ->  Tuple:
        with torch.no_grad():
         lowercase__				:    List[str]  =					EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D''' )
         model.eval()

         lowercase__				:    Optional[Any]  =					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
         lowercase__				:    str  =					model(__lowerCAmelCase )[0]
         # compare the actual values for a slice.
         lowercase__				:    Optional[Any]  =					torch.tensor(
             [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]] )
         self.assertTrue(torch.allclose(output[:, :3, :3]	,   __lowerCAmelCase	,   atol=1E-4 ) )


'''simple docstring'''

from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import convert_to_rgb, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
  import PIL


__a: List[str]          =   logging.get_logger(__name__)



class 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE							=    ["pixel_values"]

       def __init__(						self	,   __lowerCAmelCase = True	,   __lowerCAmelCase = None	,   __lowerCAmelCase = PILImageResampling.BICUBIC	,   __lowerCAmelCase = True	,   __lowerCAmelCase = 1 / 255	,   __lowerCAmelCase = True	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = True	,   **__lowerCAmelCase	,   )  ->  None:
        super().__init__(**__lowerCAmelCase )
        lowercase__				:    Optional[int]  =					size if size is not None else {'''height''': 384, '''width''': 384}
        lowercase__				:    Optional[Any]  =					get_size_dict(__lowerCAmelCase	,   default_to_square=__lowerCAmelCase )

        lowercase__				:    Dict  =					do_resize
        lowercase__				:    int  =					size
        lowercase__				:    int  =					resample
        lowercase__				:    Tuple  =					do_rescale
        lowercase__				:    int  =					rescale_factor
        lowercase__				:    int  =					do_normalize
        lowercase__				:    Optional[int]  =					image_mean if image_mean is not None else OPENAI_CLIP_MEAN
        lowercase__				:    Optional[int]  =					image_std if image_std is not None else OPENAI_CLIP_STD
        lowercase__				:    Tuple  =					do_convert_rgb

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = PILImageResampling.BICUBIC	,   __lowerCAmelCase = None	,   **__lowerCAmelCase	,   )  ->  np.ndarray:
        lowercase__				:    Union[str, Any]  =					get_size_dict(__lowerCAmelCase	,   default_to_square=__lowerCAmelCase )
        if "height" not in size or "width" not in size:
         raise ValueError(F"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}""" )
        lowercase__				:    Any  =					(size['''height'''], size['''width'''])
        return resize(__lowerCAmelCase	,   size=__lowerCAmelCase	,   resample=__lowerCAmelCase	,   data_format=__lowerCAmelCase	,   **__lowerCAmelCase )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = None	,   **__lowerCAmelCase	,   )  ->  str:
        return rescale(__lowerCAmelCase	,   scale=__lowerCAmelCase	,   data_format=__lowerCAmelCase	,   **__lowerCAmelCase )

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = None	,   **__lowerCAmelCase	,   )  ->  np.ndarray:
        return normalize(__lowerCAmelCase	,   mean=__lowerCAmelCase	,   std=__lowerCAmelCase	,   data_format=__lowerCAmelCase	,   **__lowerCAmelCase )


       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = ChannelDimension.FIRST	,   **__lowerCAmelCase	,   )  ->  PIL.Image.Image:
        lowercase__				:    Union[str, Any]  =					do_resize if do_resize is not None else self.do_resize
        lowercase__				:    Any  =					resample if resample is not None else self.resample
        lowercase__				:    int  =					do_rescale if do_rescale is not None else self.do_rescale
        lowercase__				:    List[Any]  =					rescale_factor if rescale_factor is not None else self.rescale_factor
        lowercase__				:    List[Any]  =					do_normalize if do_normalize is not None else self.do_normalize
        lowercase__				:    Dict  =					image_mean if image_mean is not None else self.image_mean
        lowercase__				:    Dict  =					image_std if image_std is not None else self.image_std
        lowercase__				:    Dict  =					do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

        lowercase__				:    Optional[int]  =					size if size is not None else self.size
        lowercase__				:    int  =					get_size_dict(__lowerCAmelCase	,   default_to_square=__lowerCAmelCase )

        lowercase__				:    str  =					make_list_of_images(__lowerCAmelCase )

        if not valid_images(__lowerCAmelCase ):
         raise ValueError(
             '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
             '''torch.Tensor, tf.Tensor or jax.ndarray.''' )

        if do_resize and size is None or resample is None:
         raise ValueError('''Size and resample must be specified if do_resize is True.''' )

        if do_rescale and rescale_factor is None:
         raise ValueError('''Rescale factor must be specified if do_rescale is True.''' )

        if do_normalize and (image_mean is None or image_std is None):
         raise ValueError('''Image mean and std must be specified if do_normalize is True.''' )

        # PIL RGBA images are converted to RGB
        if do_convert_rgb:
         lowercase__				:    Optional[Any]  =					[convert_to_rgb(__lowerCAmelCase ) for image in images]

        # All transformations expect numpy arrays.
        lowercase__				:    Any  =					[to_numpy_array(__lowerCAmelCase ) for image in images]

        if do_resize:
         lowercase__				:    Tuple  =					[self.resize(image=__lowerCAmelCase	,   size=__lowerCAmelCase	,   resample=__lowerCAmelCase ) for image in images]

        if do_rescale:
         lowercase__				:    List[str]  =					[self.rescale(image=__lowerCAmelCase	,   scale=__lowerCAmelCase ) for image in images]

        if do_normalize:
         lowercase__				:    Tuple  =					[self.normalize(image=__lowerCAmelCase	,   mean=__lowerCAmelCase	,   std=__lowerCAmelCase ) for image in images]

        lowercase__				:    List[str]  =					[to_channel_dimension_format(__lowerCAmelCase	,   __lowerCAmelCase ) for image in images]

        lowercase__				:    Optional[Any]  =					BatchFeature(data={'''pixel_values''': images}	,   tensor_type=__lowerCAmelCase )

        return encoded_outputs

