



import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor, SwinConfig
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCAmelCase	       =	logging.get_logger(__name__)
def 		UpperCamelCase     (       snake_case__	:   Optional[Any]     )     ->					Optional[int]:
   UpperCamelCase   :  List[str]						     =     SwinConfig(
       embed_dim=192	,					depths=(2, 2, 18, 2)	,					num_heads=(6, 12, 24, 48)	,					window_size=12	,					out_features=['stage2', 'stage3', 'stage4']	,					)

   UpperCamelCase   :  List[str]						     =     DetaConfig(
       backbone_config=snake_case__	,					num_queries=900	,					encoder_ffn_dim=2048	,					decoder_ffn_dim=2048	,					num_feature_levels=5	,					assign_first_stage=snake_case__	,					with_box_refine=snake_case__	,					two_stage=snake_case__	,					)

   # set labels
   UpperCamelCase   :  Dict						     =     'huggingface/label-files'
   if "o365" in model_name:
      UpperCamelCase   :  Union[str, Any]						     =     366
      UpperCamelCase   :  Optional[int]						     =     'object365-id2label.json'
   else:
      UpperCamelCase   :  Tuple						     =     91
      UpperCamelCase   :  Optional[int]						     =     'coco-detection-id2label.json'

   UpperCamelCase   :  List[Any]						     =     num_labels
   UpperCamelCase   :  Any						     =     json.load(open(cached_download(hf_hub_url(snake_case__	,					snake_case__	,					repo_type='dataset'     )     )	,					'r'     )     )
   UpperCamelCase   :  Optional[int]						     =     {int(snake_case__     ): v for k, v in idalabel.items()}
   UpperCamelCase   :  List[str]						     =     idalabel
   UpperCamelCase   :  Any						     =     {v: k for k, v in idalabel.items()}

   return config
def 		UpperCamelCase     (       snake_case__	:   Any     )     ->					Any:
   UpperCamelCase   :  List[str]						     =     []

   # stem
   # fmt: off
   rename_keys.append(('backbone.0.body.patch_embed.proj.weight', 'model.backbone.model.embeddings.patch_embeddings.projection.weight')     )
   rename_keys.append(('backbone.0.body.patch_embed.proj.bias', 'model.backbone.model.embeddings.patch_embeddings.projection.bias')     )
   rename_keys.append(('backbone.0.body.patch_embed.norm.weight', 'model.backbone.model.embeddings.norm.weight')     )
   rename_keys.append(('backbone.0.body.patch_embed.norm.bias', 'model.backbone.model.embeddings.norm.bias')     )
   # stages
   for i in range(len(config.backbone_config.depths     )     ):
      for j in range(config.backbone_config.depths[i]     ):
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.norm1.weight""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.norm1.bias""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_bias_table""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_index""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.attn.proj.weight""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.attn.proj.bias""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.norm2.weight""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.norm2.bias""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.weight""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.bias""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.weight""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.weight""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.bias""", F"""model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.bias""")     )

      if i < 3:
         rename_keys.append((F"""backbone.0.body.layers.{i}.downsample.reduction.weight""", F"""model.backbone.model.encoder.layers.{i}.downsample.reduction.weight""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.downsample.norm.weight""", F"""model.backbone.model.encoder.layers.{i}.downsample.norm.weight""")     )
         rename_keys.append((F"""backbone.0.body.layers.{i}.downsample.norm.bias""", F"""model.backbone.model.encoder.layers.{i}.downsample.norm.bias""")     )

   rename_keys.append(('backbone.0.body.norm1.weight', 'model.backbone.model.hidden_states_norms.stage2.weight')     )
   rename_keys.append(('backbone.0.body.norm1.bias', 'model.backbone.model.hidden_states_norms.stage2.bias')     )
   rename_keys.append(('backbone.0.body.norm2.weight', 'model.backbone.model.hidden_states_norms.stage3.weight')     )
   rename_keys.append(('backbone.0.body.norm2.bias', 'model.backbone.model.hidden_states_norms.stage3.bias')     )
   rename_keys.append(('backbone.0.body.norm3.weight', 'model.backbone.model.hidden_states_norms.stage4.weight')     )
   rename_keys.append(('backbone.0.body.norm3.bias', 'model.backbone.model.hidden_states_norms.stage4.bias')     )

   # transformer encoder
   for i in range(config.encoder_layers     ):
      rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight""", F"""model.encoder.layers.{i}.self_attn.sampling_offsets.weight""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias""", F"""model.encoder.layers.{i}.self_attn.sampling_offsets.bias""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.attention_weights.weight""", F"""model.encoder.layers.{i}.self_attn.attention_weights.weight""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.attention_weights.bias""", F"""model.encoder.layers.{i}.self_attn.attention_weights.bias""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.value_proj.weight""", F"""model.encoder.layers.{i}.self_attn.value_proj.weight""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.value_proj.bias""", F"""model.encoder.layers.{i}.self_attn.value_proj.bias""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.output_proj.weight""", F"""model.encoder.layers.{i}.self_attn.output_proj.weight""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.output_proj.bias""", F"""model.encoder.layers.{i}.self_attn.output_proj.bias""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.norm1.weight""", F"""model.encoder.layers.{i}.self_attn_layer_norm.weight""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.norm1.bias""", F"""model.encoder.layers.{i}.self_attn_layer_norm.bias""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.weight""", F"""model.encoder.layers.{i}.fc1.weight""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.bias""", F"""model.encoder.layers.{i}.fc1.bias""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.weight""", F"""model.encoder.layers.{i}.fc2.weight""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.bias""", F"""model.encoder.layers.{i}.fc2.bias""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.weight""", F"""model.encoder.layers.{i}.final_layer_norm.weight""")     )
      rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.bias""", F"""model.encoder.layers.{i}.final_layer_norm.bias""")     )

   # transformer decoder
   for i in range(config.decoder_layers     ):
      rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight""", F"""model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias""", F"""model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.attention_weights.weight""", F"""model.decoder.layers.{i}.encoder_attn.attention_weights.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.attention_weights.bias""", F"""model.decoder.layers.{i}.encoder_attn.attention_weights.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.value_proj.weight""", F"""model.decoder.layers.{i}.encoder_attn.value_proj.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.value_proj.bias""", F"""model.decoder.layers.{i}.encoder_attn.value_proj.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.output_proj.weight""", F"""model.decoder.layers.{i}.encoder_attn.output_proj.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.output_proj.bias""", F"""model.decoder.layers.{i}.encoder_attn.output_proj.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.norm1.weight""", F"""model.decoder.layers.{i}.encoder_attn_layer_norm.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.norm1.bias""", F"""model.decoder.layers.{i}.encoder_attn_layer_norm.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.self_attn.out_proj.weight""", F"""model.decoder.layers.{i}.self_attn.out_proj.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.self_attn.out_proj.bias""", F"""model.decoder.layers.{i}.self_attn.out_proj.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.norm2.weight""", F"""model.decoder.layers.{i}.self_attn_layer_norm.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.norm2.bias""", F"""model.decoder.layers.{i}.self_attn_layer_norm.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.weight""", F"""model.decoder.layers.{i}.fc1.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.bias""", F"""model.decoder.layers.{i}.fc1.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.weight""", F"""model.decoder.layers.{i}.fc2.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.bias""", F"""model.decoder.layers.{i}.fc2.bias""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.weight""", F"""model.decoder.layers.{i}.final_layer_norm.weight""")     )
      rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.bias""", F"""model.decoder.layers.{i}.final_layer_norm.bias""")     )

   # fmt: on

   return rename_keys
def 		UpperCamelCase     (       snake_case__	:   Tuple	,					snake_case__	:   Optional[int]	,					snake_case__	:   str     )     ->					Optional[int]:
   UpperCamelCase   :  Tuple						     =     dct.pop(snake_case__     )
   UpperCamelCase   :  Any						     =     val
def 		UpperCamelCase     (       snake_case__	:   int	,					snake_case__	:   str     )     ->					Union[str, Any]:
   UpperCamelCase   :  Any						     =     [int(backbone_config.embed_dim * 2**i     ) for i in range(len(backbone_config.depths     )     )]
   for i in range(len(backbone_config.depths     )     ):
      UpperCamelCase   :  Union[str, Any]						     =     num_features[i]
      for j in range(backbone_config.depths[i]     ):
         # fmt: off
         # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
         UpperCamelCase   :  Optional[Any]						     =     state_dict.pop(F"""backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight"""     )
         UpperCamelCase   :  List[Any]						     =     state_dict.pop(F"""backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.bias"""     )
         # next, add query, keys and values (in that order) to the state dict
         UpperCamelCase   :  Optional[Any]						     =     in_proj_weight[:dim, :]
         UpperCamelCase   :  int						     =     in_proj_bias[: dim]
         UpperCamelCase   :  Tuple						     =     in_proj_weight[
             dim : dim * 2, :
         ]
         UpperCamelCase   :  List[str]						     =     in_proj_bias[
             dim : dim * 2
         ]
         UpperCamelCase   :  List[Any]						     =     in_proj_weight[
             -dim :, :
         ]
         UpperCamelCase   :  List[Any]						     =     in_proj_bias[-dim :]
         # fmt: on
def 		UpperCamelCase     (       snake_case__	:   Optional[Any]	,					snake_case__	:   int     )     ->					Optional[Any]:
   # transformer decoder self-attention layers
   UpperCamelCase   :  Tuple						     =     config.d_model
   for i in range(config.decoder_layers     ):
      # read in weights + bias of input projection layer of self-attention
      UpperCamelCase   :  List[Any]						     =     state_dict.pop(F"""transformer.decoder.layers.{i}.self_attn.in_proj_weight"""     )
      UpperCamelCase   :  Optional[int]						     =     state_dict.pop(F"""transformer.decoder.layers.{i}.self_attn.in_proj_bias"""     )
      # next, add query, keys and values (in that order) to the state dict
      UpperCamelCase   :  Optional[Any]						     =     in_proj_weight[:hidden_size, :]
      UpperCamelCase   :  Any						     =     in_proj_bias[:hidden_size]
      UpperCamelCase   :  Union[str, Any]						     =     in_proj_weight[
          hidden_size : hidden_size * 2, :
      ]
      UpperCamelCase   :  Dict						     =     in_proj_bias[hidden_size : hidden_size * 2]
      UpperCamelCase   :  Tuple						     =     in_proj_weight[-hidden_size:, :]
      UpperCamelCase   :  List[Any]						     =     in_proj_bias[-hidden_size:]
def 		UpperCamelCase     (       )     ->					Optional[int]:
   UpperCamelCase   :  int						     =     'http://images.cocodataset.org/val2017/000000039769.jpg'
   UpperCamelCase   :  int						     =     Image.open(requests.get(snake_case__	,					stream=snake_case__     ).raw     )

   return im
@torch.no_grad()
def 		UpperCamelCase     (       snake_case__	:   List[Any]	,					snake_case__	:   Dict	,					snake_case__	:   Tuple     )     ->					str:
   UpperCamelCase   :  Any						     =     get_deta_config(snake_case__     )

   # load original state dict
   if model_name == "deta-swin-large":
      UpperCamelCase   :  Optional[Any]						     =     hf_hub_download(repo_id='nielsr/deta-checkpoints'	,					filename='adet_swin_ft.pth'     )
   elif model_name == "deta-swin-large-o365":
      UpperCamelCase   :  List[Any]						     =     hf_hub_download(repo_id='jozhang97/deta-swin-l-o365'	,					filename='deta_swin_pt_o365.pth'     )
   else:
      raise ValueError(F"""Model name {model_name} not supported"""     )

   UpperCamelCase   :  Union[str, Any]						     =     torch.load(snake_case__	,					map_location='cpu'     )['model']

   # original state dict
   for name, param in state_dict.items():
      print(snake_case__	,					param.shape     )

   # rename keys
   UpperCamelCase   :  List[str]						     =     create_rename_keys(snake_case__     )
   for src, dest in rename_keys:
      rename_key(snake_case__	,					snake_case__	,					snake_case__     )
   read_in_swin_q_k_v(snake_case__	,					config.backbone_config     )
   read_in_decoder_q_k_v(snake_case__	,					snake_case__     )

   # fix some prefixes
   for key in state_dict.copy().keys():
      if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
         UpperCamelCase   :  Dict						     =     state_dict.pop(snake_case__     )
         UpperCamelCase   :  Any						     =     val
      if "input_proj" in key:
         UpperCamelCase   :  Optional[Any]						     =     state_dict.pop(snake_case__     )
         UpperCamelCase   :  Tuple						     =     val
      if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
         UpperCamelCase   :  str						     =     state_dict.pop(snake_case__     )
         UpperCamelCase   :  str						     =     val

    # finally, create HuggingFace model and load state dict
   UpperCamelCase   :  Tuple						     =     DetaForObjectDetection(snake_case__     )
   model.load_state_dict(snake_case__     )
   model.eval()

   UpperCamelCase   :  Optional[int]						     =     'cuda' if torch.cuda.is_available() else 'cpu'
   model.to(snake_case__     )

   # load image processor
   UpperCamelCase   :  Union[str, Any]						     =     DetaImageProcessor(format='coco_detection'     )

   # verify our conversion on image
   UpperCamelCase   :  str						     =     prepare_img()
   UpperCamelCase   :  Any						     =     processor(images=snake_case__	,					return_tensors='pt'     )
   UpperCamelCase   :  Any						     =     encoding['pixel_values']
   UpperCamelCase   :  Any						     =     model(pixel_values.to(snake_case__     )     )

   # verify logits
   print('Logits:'	,					outputs.logits[0, :3, :3]     )
   print('Boxes:'	,					outputs.pred_boxes[0, :3, :3]     )
   if model_name == "deta-swin-large":
      UpperCamelCase   :  Dict						     =     torch.tensor(
          [[-7.6308, -2.8485, -5.3737], [-7.2037, -4.5505, -4.8027], [-7.2943, -4.2611, -4.6617]]     )
      UpperCamelCase   :  str						     =     torch.tensor([[0.4987, 0.4969, 0.9999], [0.2549, 0.5498, 0.4805], [0.5498, 0.2757, 0.0569]]     )
   elif model_name == "deta-swin-large-o365":
      UpperCamelCase   :  Union[str, Any]						     =     torch.tensor(
          [[-8.0122, -3.5720, -4.9717], [-8.1547, -3.6886, -4.6389], [-7.6610, -3.6194, -5.0134]]     )
      UpperCamelCase   :  Union[str, Any]						     =     torch.tensor([[0.2523, 0.5549, 0.4881], [0.7715, 0.4149, 0.4601], [0.5503, 0.2753, 0.0575]]     )
   assert torch.allclose(outputs.logits[0, :3, :3]	,					expected_logits.to(snake_case__     )	,					atol=1E-4     )
   assert torch.allclose(outputs.pred_boxes[0, :3, :3]	,					expected_boxes.to(snake_case__     )	,					atol=1E-4     )
   print('Everything ok!'     )

   if pytorch_dump_folder_path:
      # Save model and processor
      logger.info(F"""Saving PyTorch model and processor to {pytorch_dump_folder_path}..."""     )
      Path(snake_case__     ).mkdir(exist_ok=snake_case__     )
      model.save_pretrained(snake_case__     )
      processor.save_pretrained(snake_case__     )

   # Push to hub
   if push_to_hub:
      print('Pushing model and processor to hub...'     )
      model.push_to_hub(F"""jozhang97/{model_name}"""     )
      processor.push_to_hub(F"""jozhang97/{model_name}"""     )


if __name__ == "__main__":
    __UpperCAmelCase	       =	argparse.ArgumentParser()

    parser.add_argument(
        '''--model_name''',
        type=str,
        default='''deta-swin-large''',
        choices=['''deta-swin-large''', '''deta-swin-large-o365'''],
        help='''Name of the model you\'d like to convert.''',
    )
    parser.add_argument(
        '''--pytorch_dump_folder_path''',
        default=None,
        type=str,
        help='''Path to the folder to output PyTorch model.''',
    )
    parser.add_argument(
        '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
    )
    __UpperCAmelCase	       =	parser.parse_args()
    convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








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     lowerCAmelCase_		:
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_=13,      SCREAMING_SNAKE_CASE_=7,      SCREAMING_SNAKE_CASE_=False,      SCREAMING_SNAKE_CASE_=True,      SCREAMING_SNAKE_CASE_=False,      SCREAMING_SNAKE_CASE_=False,      SCREAMING_SNAKE_CASE_=19,      SCREAMING_SNAKE_CASE_=32,      SCREAMING_SNAKE_CASE_=5,      SCREAMING_SNAKE_CASE_=4,      SCREAMING_SNAKE_CASE_=37,      SCREAMING_SNAKE_CASE_="gelu",      SCREAMING_SNAKE_CASE_=0.1,      SCREAMING_SNAKE_CASE_=0.1,      SCREAMING_SNAKE_CASE_=512,      SCREAMING_SNAKE_CASE_=16,      SCREAMING_SNAKE_CASE_=2,      SCREAMING_SNAKE_CASE_=0.02,      SCREAMING_SNAKE_CASE_=3,      SCREAMING_SNAKE_CASE_=4,      SCREAMING_SNAKE_CASE_=None,      )					->		Optional[int]:
          UpperCamelCase   :  Any						     =     parent
          UpperCamelCase   :  Optional[Any]						     =     batch_size
          UpperCamelCase   :  Optional[Any]						     =     seq_length
          UpperCamelCase   :  Tuple						     =     is_training
          UpperCamelCase   :  List[str]						     =     use_input_mask
          UpperCamelCase   :  Optional[int]						     =     use_token_type_ids
          UpperCamelCase   :  Any						     =     use_labels
          UpperCamelCase   :  Union[str, Any]						     =     vocab_size
          UpperCamelCase   :  int						     =     hidden_size
          UpperCamelCase   :  List[Any]						     =     num_hidden_layers
          UpperCamelCase   :  Dict						     =     num_attention_heads
          UpperCamelCase   :  List[Any]						     =     intermediate_size
          UpperCamelCase   :  str						     =     hidden_act
          UpperCamelCase   :  List[Any]						     =     hidden_dropout_prob
          UpperCamelCase   :  List[str]						     =     attention_probs_dropout_prob
          UpperCamelCase   :  List[str]						     =     max_position_embeddings
          UpperCamelCase   :  Union[str, Any]						     =     type_vocab_size
          UpperCamelCase   :  str						     =     type_sequence_label_size
          UpperCamelCase   :  List[str]						     =     initializer_range
          UpperCamelCase   :  List[Any]						     =     num_labels
          UpperCamelCase   :  Any						     =     num_choices
          UpperCamelCase   :  Tuple						     =     scope
       def  snake_case_	(      self	)					->		Optional[int]:
          UpperCamelCase   :  str						     =     ids_tensor([self.batch_size, self.seq_length],      self.vocab_size	)

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

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

          UpperCamelCase   :  Union[str, Any]						     =     self.get_config()

          return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
       def  snake_case_	(      self	)					->		Optional[int]:
          UpperCamelCase   :  List[str]						     =     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=SCREAMING_SNAKE_CASE_,      esmfold_config={'trunk': {'num_blocks': 2}, 'fp16_esm': False},      )
          return config
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)					->		Tuple:
          UpperCamelCase   :  Union[str, Any]						     =     EsmForProteinFolding(config=SCREAMING_SNAKE_CASE_	).float()
          model.to(SCREAMING_SNAKE_CASE_	)
          model.eval()
          UpperCamelCase   :  int						     =     model(SCREAMING_SNAKE_CASE_,      attention_mask=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  List[Any]						     =     model(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  List[Any]						     =     model(SCREAMING_SNAKE_CASE_	)

          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  snake_case_	(      self	)					->		Optional[int]:
          UpperCamelCase   :  Union[str, Any]						     =     self.prepare_config_and_inputs()
          (
              (
              UpperCamelCase
          )     ,					(
              UpperCamelCase
          )     ,					(
              UpperCamelCase
          )     ,					(
              UpperCamelCase
          )     ,					(
              UpperCamelCase
          )     ,					(
              UpperCamelCase
          )     ,					
          )   :  Tuple						     =     config_and_inputs
          UpperCamelCase   :  Optional[Any]						     =     {'input_ids': input_ids, 'attention_mask': input_mask}
          return config, inputs_dict






@require_torch
class     lowerCAmelCase_		(  a__	,   a__	,   unittest.TestCase							):
       UpperCAmelCase__    : Dict       					=  False

       UpperCAmelCase__    : Optional[int]       					=  (EsmForProteinFolding,) if is_torch_available() else ()
       UpperCAmelCase__    : int       					=  ()
       UpperCAmelCase__    : List[str]       					=  {} if is_torch_available() else {}
       UpperCAmelCase__    : Optional[int]       					=  False
       def  snake_case_	(      self	)					->		Dict:
          UpperCamelCase   :  Tuple						     =     EsmFoldModelTester(self	)
          UpperCamelCase   :  List[str]						     =     ConfigTester(self,      config_class=SCREAMING_SNAKE_CASE_,      hidden_size=37	)
       def  snake_case_	(      self	)					->		int:
          self.config_tester.run_common_tests()
       def  snake_case_	(      self	)					->		Union[str, Any]:
          UpperCamelCase   :  int						     =     self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)
       @unittest.skip('Does not support attention outputs'	)
       def  snake_case_	(      self	)					->		Tuple:
          pass
       @unittest.skip
       def  snake_case_	(      self	)					->		List[Any]:
          pass
       @unittest.skip('Esm does not support embedding resizing'	)
       def  snake_case_	(      self	)					->		Any:
          pass
       @unittest.skip('Esm does not support embedding resizing'	)
       def  snake_case_	(      self	)					->		Optional[Any]:
          pass
       @unittest.skip('ESMFold does not support passing input embeds!'	)
       def  snake_case_	(      self	)					->		Optional[Any]:
          pass
       @unittest.skip('ESMFold does not support head pruning.'	)
       def  snake_case_	(      self	)					->		Any:
          pass
       @unittest.skip('ESMFold does not support head pruning.'	)
       def  snake_case_	(      self	)					->		int:
          pass
       @unittest.skip('ESMFold does not support head pruning.'	)
       def  snake_case_	(      self	)					->		Dict:
          pass
       @unittest.skip('ESMFold does not support head pruning.'	)
       def  snake_case_	(      self	)					->		Union[str, Any]:
          pass
       @unittest.skip('ESMFold does not support head pruning.'	)
       def  snake_case_	(      self	)					->		Any:
          pass
       @unittest.skip('ESMFold does not output hidden states in the normal way.'	)
       def  snake_case_	(      self	)					->		str:
          pass
       @unittest.skip('ESMfold does not output hidden states in the normal way.'	)
       def  snake_case_	(      self	)					->		List[str]:
          pass
       @unittest.skip('ESMFold only has one output format.'	)
       def  snake_case_	(      self	)					->		int:
          pass
       @unittest.skip('This test doesn\'t work for ESMFold and doesn\'t test core functionality'	)
       def  snake_case_	(      self	)					->		Any:
          pass
       @unittest.skip('ESMFold does not support input chunking.'	)
       def  snake_case_	(      self	)					->		Any:
          pass
       @unittest.skip('ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'	)
       def  snake_case_	(      self	)					->		Tuple:
          pass
       @unittest.skip('ESMFold doesn\'t support torchscript compilation.'	)
       def  snake_case_	(      self	)					->		Any:
          pass
       @unittest.skip('ESMFold doesn\'t support torchscript compilation.'	)
       def  snake_case_	(      self	)					->		str:
          pass
       @unittest.skip('ESMFold doesn\'t support torchscript compilation.'	)
       def  snake_case_	(      self	)					->		List[Any]:
          pass
       @unittest.skip('ESMFold doesn\'t support data parallel.'	)
       def  snake_case_	(      self	)					->		List[str]:
          pass






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






@require_torch
class     lowerCAmelCase_		(  a__							):
       @slow
       def  snake_case_	(      self	)					->		str:
          UpperCamelCase   :  Union[str, Any]						     =     EsmForProteinFolding.from_pretrained('facebook/esmfold_v1'	).float()
          model.eval()
          UpperCamelCase   :  Tuple						     =     torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]	)
          UpperCamelCase   :  List[str]						     =     model(SCREAMING_SNAKE_CASE_	)['positions']
          UpperCamelCase   :  int						     =     torch.tensor([2.58_28, 0.79_93, -10.93_34],      dtype=torch.floataa	)
          self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0],      SCREAMING_SNAKE_CASE_,      atol=1e-4	)	)



