
import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import UperNetForSemanticSegmentation
 from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor


class 			lowercase_		:


       def __init__(    self						,       __A						,       __A=13						,       __A=32						,       __A=3						,       __A=4						,       __A=[10, 20, 30, 40]						,       __A=[2, 2, 3, 2]						,       __A=True						,       __A=True						,       __A=37						,       __A="gelu"						,       __A=10						,       __A=0.02						,       __A=["stage2", "stage3", "stage4"]						,       __A=3						,       __A=None						,       )						->    int:
              SCREAMING_SNAKE_CASE_							:				str		=parent
              SCREAMING_SNAKE_CASE_							:				List[Any]		=batch_size
              SCREAMING_SNAKE_CASE_							:				List[Any]		=image_size
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=num_channels
              SCREAMING_SNAKE_CASE_							:				str		=num_stages
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=hidden_sizes
              SCREAMING_SNAKE_CASE_							:				Tuple		=depths
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=is_training
              SCREAMING_SNAKE_CASE_							:				Dict		=use_labels
              SCREAMING_SNAKE_CASE_							:				str		=intermediate_size
              SCREAMING_SNAKE_CASE_							:				List[str]		=hidden_act
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=type_sequence_label_size
              SCREAMING_SNAKE_CASE_							:				Tuple		=initializer_range
              SCREAMING_SNAKE_CASE_							:				List[str]		=out_features
              SCREAMING_SNAKE_CASE_							:				Dict		=num_labels
              SCREAMING_SNAKE_CASE_							:				List[Any]		=scope
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=num_stages


       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

              SCREAMING_SNAKE_CASE_							:				Tuple		=None
              if self.use_labels:
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=ids_tensor([self.batch_size]						,       self.type_sequence_label_size				)

              SCREAMING_SNAKE_CASE_							:				Any		=self.get_config()

              return config, pixel_values, labels


       def 	_snake_case				(    self				)						->    List[Any]:
              return ConvNextConfig(
                  num_channels=self.num_channels						,       num_stages=self.num_stages						,       hidden_sizes=self.hidden_sizes						,       depths=self.depths						,       is_training=self.is_training						,       intermediate_size=self.intermediate_size						,       hidden_act=self.hidden_act						,       out_features=self.out_features						,       )


       def 	_snake_case				(    self				)						->    Any:
              return UperNetConfig(
                  backbone_config=self.get_backbone_config()						,       hidden_size=512						,       pool_scales=[1, 2, 3, 6]						,       use_auxiliary_head=__A						,       auxiliary_loss_weight=0.4						,       auxiliary_in_channels=40						,       auxiliary_channels=256						,       auxiliary_num_convs=1						,       auxiliary_concat_input=__A						,       loss_ignore_index=255						,       num_labels=self.num_labels						,       )


       def 	_snake_case				(    self						,       __A						,       __A						,       __A				)						->    Any:
              SCREAMING_SNAKE_CASE_							:				List[str]		=UperNetForSemanticSegmentation(config=__A				)
              model.to(__A				)
              model.eval()
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=model(__A				)
              self.parent.assertEqual(
                  result.logits.shape						,       (self.batch_size, self.num_labels, self.image_size, self.image_size)				)





       def 	_snake_case				(    self				)						->    Tuple:
              SCREAMING_SNAKE_CASE_							:				int		=self.prepare_config_and_inputs()
              (
                  (
                  SCREAMING_SNAKE_CASE_
              )	,     (
                  SCREAMING_SNAKE_CASE_
              )	,     (
                  SCREAMING_SNAKE_CASE_
              )	,     
              )							:				List[str]		=config_and_inputs
              SCREAMING_SNAKE_CASE_							:				List[str]		={'''pixel_values''': pixel_values}
              return config, inputs_dict






@require_torch
class 			lowercase_		(      A  ,      A  ,      unittest.TestCase       ):
       __lowerCamelCase				       =						(UperNetForSemanticSegmentation,) if is_torch_available() else ()
       __lowerCamelCase				       =						{"image-segmentation": UperNetForSemanticSegmentation} if is_torch_available() else {}
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False


       def 	_snake_case				(    self				)						->    Optional[Any]:
              SCREAMING_SNAKE_CASE_							:				Tuple		=UperNetModelTester(self				)
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=ConfigTester(self						,       config_class=__A						,       has_text_modality=__A						,       hidden_size=37				)


       def 	_snake_case				(    self				)						->    List[str]:
              self.create_and_test_config_common_properties()
              self.config_tester.create_and_test_config_to_json_string()
              self.config_tester.create_and_test_config_to_json_file()
              self.config_tester.create_and_test_config_from_and_save_pretrained()
              self.config_tester.create_and_test_config_with_num_labels()
              self.config_tester.check_config_can_be_init_without_params()
              self.config_tester.check_config_arguments_init()


       def 	_snake_case				(    self				)						->    Tuple:
              return


       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				int		=self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     SCREAMING_SNAKE_CASE_							:				int		=model_class(__A				)
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=inspect.signature(model.forward				)
                     # signature.parameters is an OrderedDict => so arg_names order is deterministic
                     SCREAMING_SNAKE_CASE_							:				str		=[*signature.parameters.keys()]

                     SCREAMING_SNAKE_CASE_							:				Any		=['''pixel_values''']
                     self.assertListEqual(arg_names[:1]						,       __A				)


       def 	_snake_case				(    self				)						->    Union[str, Any]:
              SCREAMING_SNAKE_CASE_							:				Any		=self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_semantic_segmentation(*__A				)


       @unittest.skip(reason='''UperNet does not use inputs_embeds'''				)
       def 	_snake_case				(    self				)						->    Tuple:
              pass


       @unittest.skip(reason='''UperNet does not support input and output embeddings'''				)
       def 	_snake_case				(    self				)						->    Optional[int]:
              pass


       @unittest.skip(reason='''UperNet does not have a base model'''				)
       def 	_snake_case				(    self				)						->    List[Any]:
              pass


       @unittest.skip(reason='''UperNet does not have a base model'''				)
       def 	_snake_case				(    self				)						->    Any:
              pass


       @require_torch_multi_gpu
       @unittest.skip(reason='''UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''				)
       def 	_snake_case				(    self				)						->    Optional[Any]:
              pass


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


       def 	_snake_case				(    self				)						->    Dict:
              def check_hidden_states_output(__A						,       __A						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				List[str]		=model_class(__A				)
                     model.to(__A				)
                     model.eval()

                     with torch.no_grad():
                            SCREAMING_SNAKE_CASE_							:				Optional[Any]		=model(**self._prepare_for_class(__A						,       __A				)				)

                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                     SCREAMING_SNAKE_CASE_							:				List[Any]		=self.model_tester.num_stages
                     self.assertEqual(len(__A				)						,       expected_num_stages + 1				)

                     # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
                     self.assertListEqual(
                         list(hidden_states[0].shape[-2:]				)						,       [self.model_tester.image_size // 4, self.model_tester.image_size // 4]						,       )

              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Tuple		=self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=True
                     check_hidden_states_output(__A						,       __A						,       __A				)

                     # check that output_hidden_states also work using config
                     del inputs_dict["output_hidden_states"]
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=True

                     check_hidden_states_output(__A						,       __A						,       __A				)


       def 	_snake_case				(    self				)						->    Any:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				List[str]		=self.model_tester.prepare_config_and_inputs_for_common()

              SCREAMING_SNAKE_CASE_							:				Dict		=_config_zero_init(__A				)
              SCREAMING_SNAKE_CASE_							:				Any		=_config_zero_init(configs_no_init.backbone_config				)
              for model_class in self.all_model_classes:
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=model_class(config=__A				)
                     for name, param in model.named_parameters():
                            if param.requires_grad:
                                   self.assertIn(
                                       ((param.data.mean() * 1e9).round() / 1e9).item()						,       [0.0, 1.0]						,       msg=F'Parameter {name} of model {model_class} seems not properly initialized'						,       )


       @unittest.skip(reason='''UperNet does not have tied weights'''				)
       def 	_snake_case				(    self				)						->    List[str]:
              pass





       @slow
       def 	_snake_case				(    self				)						->    Optional[int]:
              for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                     SCREAMING_SNAKE_CASE_							:				Any		=UperNetForSemanticSegmentation.from_pretrained(__A				)
                     self.assertIsNotNone(__A				)






def 							SCREAMING_SNAKE_CASE_	(					)			->						List[str]:
       SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=hf_hub_download(
           repo_id='''hf-internal-testing/fixtures_ade20k'''			,							repo_type='''dataset'''			,							filename='''ADE_val_00000001.jpg'''				)
       SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=Image.open(UpperCAmelCase_				).convert('''RGB'''				)
       return image


@require_torch
@require_vision
@slow
class 			lowercase_		(      unittest.TestCase       ):


       def 	_snake_case				(    self				)						->    Tuple:
              SCREAMING_SNAKE_CASE_							:				List[str]		=AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny'''				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny'''				).to(__A				)

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=prepare_img()
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=processor(images=__A						,       return_tensors='''pt'''				).to(__A				)

              with torch.no_grad():
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=model(**__A				)

              SCREAMING_SNAKE_CASE_							:				Dict		=torch.Size((1, model.config.num_labels, 512, 512)				)
              self.assertEqual(outputs.logits.shape						,       __A				)

              SCREAMING_SNAKE_CASE_							:				Any		=torch.tensor(
                  [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]				).to(__A				)
              self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]						,       __A						,       atol=1e-4				)				)





       def 	_snake_case				(    self				)						->    str:
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny'''				)
              SCREAMING_SNAKE_CASE_							:				str		=UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny'''				).to(__A				)

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=prepare_img()
              SCREAMING_SNAKE_CASE_							:				List[Any]		=processor(images=__A						,       return_tensors='''pt'''				).to(__A				)

              with torch.no_grad():
                     SCREAMING_SNAKE_CASE_							:				List[str]		=model(**__A				)

              SCREAMING_SNAKE_CASE_							:				List[str]		=torch.Size((1, model.config.num_labels, 512, 512)				)
              self.assertEqual(outputs.logits.shape						,       __A				)

              SCREAMING_SNAKE_CASE_							:				List[str]		=torch.tensor(
                  [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]				).to(__A				)
              self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]						,       __A						,       atol=1e-4				)				)


from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowercase											=    logging.get_logger(__name__)

# TODO Update this
_lowercase											=    {
    """facebook/esm-1b""": """https://huggingface.co/facebook/esm-1b/resolve/main/config.json""",
    # See all ESM models at https://huggingface.co/models?filter=esm
}


class 			lowercase_		(      A       ):
       __lowerCamelCase				       =						"esm"


       def __init__(    self						,       __A=None						,       __A=None						,       __A=None						,       __A=768						,       __A=12						,       __A=12						,       __A=3_072						,       __A=0.1						,       __A=0.1						,       __A=1_026						,       __A=0.02						,       __A=1e-1_2						,       __A="absolute"						,       __A=True						,       __A=None						,       __A=False						,       __A=False						,       __A=None						,       __A=None						,       **__A						,       )						->    int:
              super().__init__(pad_token_id=__A						,       mask_token_id=__A						,       **__A				)

              SCREAMING_SNAKE_CASE_							:				Any		=vocab_size
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=hidden_size
              SCREAMING_SNAKE_CASE_							:				int		=num_hidden_layers
              SCREAMING_SNAKE_CASE_							:				int		=num_attention_heads
              SCREAMING_SNAKE_CASE_							:				List[str]		=intermediate_size
              SCREAMING_SNAKE_CASE_							:				int		=hidden_dropout_prob
              SCREAMING_SNAKE_CASE_							:				int		=attention_probs_dropout_prob
              SCREAMING_SNAKE_CASE_							:				List[Any]		=max_position_embeddings
              SCREAMING_SNAKE_CASE_							:				List[Any]		=initializer_range
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=layer_norm_eps
              SCREAMING_SNAKE_CASE_							:				str		=position_embedding_type
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=use_cache
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=emb_layer_norm_before
              SCREAMING_SNAKE_CASE_							:				Tuple		=token_dropout
              SCREAMING_SNAKE_CASE_							:				List[Any]		=is_folding_model
              if is_folding_model:
                     if esmfold_config is None:
                            logger.info('''No esmfold_config supplied for folding model, using default values.'''				)
                            SCREAMING_SNAKE_CASE_							:				Optional[int]		=EsmFoldConfig()
                     elif isinstance(__A						,       __A				):
                            SCREAMING_SNAKE_CASE_							:				List[Any]		=EsmFoldConfig(**__A				)
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=esmfold_config
                     if vocab_list is None:
                            logger.warning('''No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!'''				)
                            SCREAMING_SNAKE_CASE_							:				Tuple		=get_default_vocab_list()
                     else:
                            SCREAMING_SNAKE_CASE_							:				Optional[Any]		=vocab_list
              else:
                     SCREAMING_SNAKE_CASE_							:				List[str]		=None
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=None
              if self.esmfold_config is not None and getattr(self.esmfold_config						,       '''use_esm_attn_map'''						,       __A				):
                     raise ValueError('''The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!'''				)





       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				str		=super().to_dict()
              if isinstance(self.esmfold_config						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.esmfold_config.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						None
       __lowerCamelCase				       =						True
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						0

       __lowerCamelCase				       =						True
       __lowerCamelCase				       =						False

       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						None


       def 	_snake_case				(    self				)						->    Optional[Any]:
              if self.trunk is None:
                     SCREAMING_SNAKE_CASE_							:				Tuple		=TrunkConfig()
              elif isinstance(self.trunk						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				List[str]		=TrunkConfig(**self.trunk				)





       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				int		=asdict(self				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=self.trunk.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						4_8
       __lowerCamelCase				       =						1_0_2_4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						3_2
       __lowerCamelCase				       =						0
       __lowerCamelCase				       =						0
       __lowerCamelCase				       =						False
       __lowerCamelCase				       =						4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						None


       def 	_snake_case				(    self				)						->    int:
              if self.structure_module is None:
                     SCREAMING_SNAKE_CASE_							:				Dict		=StructureModuleConfig()
              elif isinstance(self.structure_module						,       __A				):
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=StructureModuleConfig(**self.structure_module				)

              if self.max_recycles <= 0:
                     raise ValueError(F'`max_recycles` should be positive, got {self.max_recycles}.'				)
              if self.sequence_state_dim % self.sequence_state_dim != 0:
                     raise ValueError(
                         '''`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got'''
                         F' {self.sequence_state_dim} and {self.sequence_state_dim}.'				)
              if self.pairwise_state_dim % self.pairwise_state_dim != 0:
                     raise ValueError(
                         '''`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got'''
                         F' {self.pairwise_state_dim} and {self.pairwise_state_dim}.'				)

              SCREAMING_SNAKE_CASE_							:				Tuple		=self.sequence_state_dim // self.sequence_head_width
              SCREAMING_SNAKE_CASE_							:				str		=self.pairwise_state_dim // self.pairwise_head_width

              if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
                     raise ValueError(
                         '''`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got'''
                         F' {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}.'				)
              if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
                     raise ValueError(
                         '''`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got'''
                         F' {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}.'				)
              if self.pairwise_state_dim % 2 != 0:
                     raise ValueError(F'`pairwise_state_dim` should be even, got {self.pairwise_state_dim}.'				)

              if self.dropout >= 0.4:
                     raise ValueError(F'`dropout` should not be greater than 0.4, got {self.dropout}.'				)





       def 	_snake_case				(    self				)						->    int:
              SCREAMING_SNAKE_CASE_							:				str		=asdict(self				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=self.structure_module.to_dict()
              return output






@dataclass
class 			lowercase_		:
       __lowerCamelCase				       =						3_8_4
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						1_6
       __lowerCamelCase				       =						1_2_8
       __lowerCamelCase				       =						1_2
       __lowerCamelCase				       =						4
       __lowerCamelCase				       =						8
       __lowerCamelCase				       =						0.1
       __lowerCamelCase				       =						8
       __lowerCamelCase				       =						1
       __lowerCamelCase				       =						2
       __lowerCamelCase				       =						7
       __lowerCamelCase				       =						1_0
       __lowerCamelCase				       =						1E-8
       __lowerCamelCase				       =						1E5


       def 	_snake_case				(    self				)						->    str:
              return asdict(self				)






def 							SCREAMING_SNAKE_CASE_	(					)			->						int:
       return (
           "<cls>",
           "<pad>",
           "<eos>",
           "<unk>",
           "L",
           "A",
           "G",
           "V",
           "S",
           "E",
           "R",
           "T",
           "I",
           "D",
           "P",
           "K",
           "Q",
           "N",
           "F",
           "Y",
           "M",
           "H",
           "W",
           "C",
           "X",
           "B",
           "U",
           "Z",
           "O",
           ".",
           "-",
           "<null_1>",
           "<mask>",
       )
