




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()
__a          =			logging.get_logger(__name__)







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Optional[int]       =      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=_lowercase			,						num_queries=900			,						encoder_ffn_dim=2048			,						decoder_ffn_dim=2048			,						num_feature_levels=5			,						assign_first_stage=_lowercase			,						with_box_refine=_lowercase			,						two_stage=_lowercase			,						)

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

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

     return config







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     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 lowerCamelCase__       (      _lowercase			,						_lowercase			,						_lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Optional[Any]       =      dct.pop(_lowercase				)
     UpperCAmelCase_      :							Union[str, Any]       =      val







def lowerCamelCase__       (      _lowercase			,						_lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Tuple       =      [int(backbone_config.embed_dim * 2**i				) for i in range(len(backbone_config.depths				)				)]
     for i in range(len(backbone_config.depths				)				):
          UpperCAmelCase_      :							Tuple       =      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_      :							List[Any]       =      state_dict.pop(f'''backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight'''				)
               UpperCAmelCase_      :							Optional[int]       =      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_      :							List[Any]       =      in_proj_bias[: dim]
               UpperCAmelCase_      :							Union[str, Any]       =      in_proj_weight[
                   dim : dim * 2, :
               ]
               UpperCAmelCase_      :							int       =      in_proj_bias[
                   dim : dim * 2
               ]
               UpperCAmelCase_      :							List[Any]       =      in_proj_weight[
                   -dim :, :
               ]
               UpperCAmelCase_      :							Tuple       =      in_proj_bias[-dim :]
               # fmt: on







def lowerCamelCase__       (      _lowercase			,						_lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Union[str, Any]       =      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[Any]       =      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[int]       =      in_proj_weight[:hidden_size, :]
          UpperCAmelCase_      :							List[Any]       =      in_proj_bias[:hidden_size]
          UpperCAmelCase_      :							Optional[int]       =      in_proj_weight[
              hidden_size : hidden_size * 2, :
          ]
          UpperCAmelCase_      :							List[Any]       =      in_proj_bias[hidden_size : hidden_size * 2]
          UpperCAmelCase_      :							Tuple       =      in_proj_weight[-hidden_size:, :]
          UpperCAmelCase_      :							Optional[Any]       =      in_proj_bias[-hidden_size:]







def lowerCamelCase__       (      ):

     '''simple docstring'''


     UpperCAmelCase_      :							Any       =      '''http://images.cocodataset.org/val2017/000000039769.jpg'''
     UpperCAmelCase_      :							List[Any]       =      Image.open(requests.get(_lowercase			,						stream=_lowercase				).raw				)

     return im







@torch.no_grad()
def lowerCamelCase__       (      _lowercase			,						_lowercase			,						_lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Union[str, Any]       =      get_deta_config(_lowercase				)

     # load original state dict
     if model_name == "deta-swin-large":
          UpperCAmelCase_      :							Union[str, Any]       =      hf_hub_download(repo_id='''nielsr/deta-checkpoints'''			,						filename='''adet_swin_ft.pth'''				)
     elif model_name == "deta-swin-large-o365":
          UpperCAmelCase_      :							Union[str, 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(_lowercase			,						map_location='''cpu'''				)['''model''']

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

     # rename keys
     UpperCAmelCase_      :							Union[str, Any]       =      create_rename_keys(_lowercase				)
     for src, dest in rename_keys:
          rename_key(_lowercase			,						_lowercase			,						_lowercase				)
     read_in_swin_q_k_v(_lowercase			,						config.backbone_config				)
     read_in_decoder_q_k_v(_lowercase			,						_lowercase				)

     # 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(_lowercase				)
               UpperCAmelCase_      :							Optional[Any]       =      val
          if "input_proj" in key:
               UpperCAmelCase_      :							List[str]       =      state_dict.pop(_lowercase				)
               UpperCAmelCase_      :							Tuple       =      val
          if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
               UpperCAmelCase_      :							Union[str, Any]       =      state_dict.pop(_lowercase				)
               UpperCAmelCase_      :							Dict       =      val

    # finally, create HuggingFace model and load state dict
     UpperCAmelCase_      :							Optional[int]       =      DetaForObjectDetection(_lowercase				)
     model.load_state_dict(_lowercase				)
     model.eval()

     UpperCAmelCase_      :							Dict       =      '''cuda''' if torch.cuda.is_available() else '''cpu'''
     model.to(_lowercase				)

     # load image processor
     UpperCAmelCase_      :							Optional[int]       =      DetaImageProcessor(format='''coco_detection'''				)

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

     # verify logits
     print('''Logits:'''			,						outputs.logits[0, :3, :3]				)
     print('''Boxes:'''			,						outputs.pred_boxes[0, :3, :3]				)
     if model_name == "deta-swin-large":
          UpperCAmelCase_      :							int       =      torch.tensor(
              [[-7.6308, -2.8485, -5.3737], [-7.2037, -4.5505, -4.8027], [-7.2943, -4.2611, -4.6617]]				)
          UpperCAmelCase_      :							Any       =      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_      :							List[Any]       =      torch.tensor(
              [[-8.0122, -3.5720, -4.9717], [-8.1547, -3.6886, -4.6389], [-7.6610, -3.6194, -5.0134]]				)
          UpperCAmelCase_      :							Optional[int]       =      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(_lowercase				)			,						atol=1E-4				)
     assert torch.allclose(outputs.pred_boxes[0, :3, :3]			,						expected_boxes.to(_lowercase				)			,						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(_lowercase				).mkdir(exist_ok=_lowercase				)
          model.save_pretrained(_lowercase				)
          processor.save_pretrained(_lowercase				)

     # 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__":
     __a          =			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.'
     )
     __a          =			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   __a:






    """simple docstring"""





    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       ,)			->		Dict:
         UpperCAmelCase_      :							Optional[Any]       =      parent
         UpperCAmelCase_      :							Dict       =      batch_size
         UpperCAmelCase_      :							Optional[int]       =      seq_length
         UpperCAmelCase_      :							Union[str, Any]       =      is_training
         UpperCAmelCase_      :							Any       =      use_input_mask
         UpperCAmelCase_      :							Tuple       =      use_token_type_ids
         UpperCAmelCase_      :							Optional[int]       =      use_labels
         UpperCAmelCase_      :							List[Any]       =      vocab_size
         UpperCAmelCase_      :							List[Any]       =      hidden_size
         UpperCAmelCase_      :							str       =      num_hidden_layers
         UpperCAmelCase_      :							Dict       =      num_attention_heads
         UpperCAmelCase_      :							Optional[Any]       =      intermediate_size
         UpperCAmelCase_      :							str       =      hidden_act
         UpperCAmelCase_      :							int       =      hidden_dropout_prob
         UpperCAmelCase_      :							List[Any]       =      attention_probs_dropout_prob
         UpperCAmelCase_      :							Optional[int]       =      max_position_embeddings
         UpperCAmelCase_      :							Optional[int]       =      type_vocab_size
         UpperCAmelCase_      :							Any       =      type_sequence_label_size
         UpperCAmelCase_      :							int       =      initializer_range
         UpperCAmelCase_      :							Any       =      num_labels
         UpperCAmelCase_      :							Optional[Any]       =      num_choices
         UpperCAmelCase_      :							List[str]       =      scope





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							List[Any]       =      ids_tensor([self.batch_size, self.seq_length]       ,self.vocab_size  )

         UpperCAmelCase_      :							str       =      None
         if self.use_input_mask:
              UpperCAmelCase_      :							Optional[int]       =      random_attention_mask([self.batch_size, self.seq_length]  )

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

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

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





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							Union[str, Any]       =      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 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )			->		int:
         UpperCAmelCase_      :							Optional[int]       =      EsmForProteinFolding(config=_SCREAMING_SNAKE_CASE  ).float()
         model.to(_SCREAMING_SNAKE_CASE  )
         model.eval()
         UpperCAmelCase_      :							Union[str, Any]       =      model(_SCREAMING_SNAKE_CASE       ,attention_mask=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      model(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							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 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							int       =      self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase_
         ),   (
             UpperCAmelCase_
         ),   (
             UpperCAmelCase_
         ),   (
             UpperCAmelCase_
         ),   (
             UpperCAmelCase_
         ),   (
             UpperCAmelCase_
         ),   
         )      :							Optional[Any]       =      config_and_inputs
         UpperCAmelCase_      :							int       =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
         return config, inputs_dict


@require_torch
class   __a(					_a    ,						_a    ,						unittest.TestCase       ):






    """simple docstring"""





    lowerCAmelCase        =				False

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





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							Optional[int]       =      EsmFoldModelTester(self  )
         UpperCAmelCase_      :							int       =      ConfigTester(self       ,config_class=_SCREAMING_SNAKE_CASE       ,hidden_size=37  )





    def 							a__				(				self  )			->		Tuple:
         self.config_tester.run_common_tests()





    def 							a__				(				self  )			->		Dict:
         UpperCAmelCase_      :							Optional[Any]       =      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 							a__				(				self  )			->		Optional[int]:
         pass





    @unittest.skip
    def 							a__				(				self  )			->		Dict:
         pass





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





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





    @unittest.skip('''ESMFold does not support passing input embeds!'''  )
    def 							a__				(				self  )			->		Any:
         pass





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





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





    @unittest.skip('''ESMFold does not support head pruning.'''  )
    def 							a__				(				self  )			->		Optional[int]:
         pass





    @unittest.skip('''ESMFold does not support head pruning.'''  )
    def 							a__				(				self  )			->		Optional[int]:
         pass





    @unittest.skip('''ESMFold does not support head pruning.'''  )
    def 							a__				(				self  )			->		Dict:
         pass





    @unittest.skip('''ESMFold does not output hidden states in the normal way.'''  )
    def 							a__				(				self  )			->		str:
         pass





    @unittest.skip('''ESMfold does not output hidden states in the normal way.'''  )
    def 							a__				(				self  )			->		Optional[Any]:
         pass





    @unittest.skip('''ESMFold only has one output format.'''  )
    def 							a__				(				self  )			->		Optional[int]:
         pass





    @unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality'''  )
    def 							a__				(				self  )			->		int:
         pass





    @unittest.skip('''ESMFold does not support input chunking.'''  )
    def 							a__				(				self  )			->		List[Any]:
         pass





    @unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''  )
    def 							a__				(				self  )			->		Tuple:
         pass





    @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''  )
    def 							a__				(				self  )			->		Optional[int]:
         pass





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





    @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''  )
    def 							a__				(				self  )			->		Tuple:
         pass





    @unittest.skip('''ESMFold doesn\'t support data parallel.'''  )
    def 							a__				(				self  )			->		Dict:
         pass





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


@require_torch
class   __a(					_a       ):






    """simple docstring"""





    @slow
    def 							a__				(				self  )			->		List[Any]:
         UpperCAmelCase_      :							List[str]       =      EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''  ).float()
         model.eval()
         UpperCAmelCase_      :							str       =      torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
         UpperCAmelCase_      :							Optional[int]       =      model(_SCREAMING_SNAKE_CASE  )['''positions''']
         UpperCAmelCase_      :							List[str]       =      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  )  )