






import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
   import torch

   from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin




class 							_lowerCAmelCase:





       """simple docstring"""







       def __init__(   self    ,     _lowerCamelCase    ,     _lowerCamelCase=None    ,     _lowerCamelCase=None    ,     _lowerCamelCase=None    ,     _lowerCamelCase="resnet50"    ,     _lowerCamelCase=3    ,     _lowerCamelCase=3_2    ,     _lowerCamelCase=3    ,     _lowerCamelCase=True    ,     _lowerCamelCase=True    ,     ):
         UpperCamelCase_:   str        =							parent
         UpperCamelCase_:   Optional[Any]        =							out_indices if out_indices is not None else [4]
         UpperCamelCase_:   int        =							stage_names
         UpperCamelCase_:   Tuple        =							out_features
         UpperCamelCase_:   Tuple        =							backbone
         UpperCamelCase_:   Union[str, Any]        =							batch_size
         UpperCamelCase_:   int        =							image_size
         UpperCamelCase_:   Dict        =							num_channels
         UpperCamelCase_:   Optional[Any]        =							use_pretrained_backbone
         UpperCamelCase_:   Optional[int]        =							is_training





       def      _a       (   self				):
         UpperCamelCase_:   Any        =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)
         UpperCamelCase_:   Optional[int]        =							self.get_config()

         return config, pixel_values





       def      _a       (   self				):
         return TimmBackboneConfig(
             image_size=self.image_size    ,     num_channels=self.num_channels    ,     out_features=self.out_features    ,     out_indices=self.out_indices    ,     stage_names=self.stage_names    ,     use_pretrained_backbone=self.use_pretrained_backbone    ,     backbone=self.backbone    ,     )





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   Union[str, Any]        =							TimmBackbone(config=_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         with torch.no_grad():
           UpperCamelCase_:   int        =							model(_lowerCamelCase				)
         self.parent.assertEqual(
             result.feature_map[-1].shape    ,     (self.batch_size, model.channels[-1], 1_4, 1_4)    ,     )





       def      _a       (   self				):
         UpperCamelCase_:   Optional[int]        =							self.prepare_config_and_inputs()
         UpperCamelCase_ ,UpperCamelCase_:   Optional[int]        =							config_and_inputs
         UpperCamelCase_:   Optional[int]        =							{'pixel_values': pixel_values}
         return config, inputs_dict







@require_torch
@require_timm
class 							_lowerCAmelCase(			UpperCAmelCase_				,      UpperCAmelCase_				,      UpperCAmelCase_				,      unittest.TestCase   ):





       """simple docstring"""







       a							:						Any										=(TimmBackbone,) if is_torch_available() else ()
       a							:						Any										={'''feature-extraction''': TimmBackbone} if is_torch_available() else {}
       a							:						Dict										=False
       a							:						Optional[int]										=False
       a							:						List[Any]										=False
       a							:						int										=False





       def      _a       (   self				):
         UpperCamelCase_:   int        =							TimmBackboneModelTester(self				)
         UpperCamelCase_:   Optional[int]        =							ConfigTester(self    ,     config_class=_lowerCamelCase    ,     has_text_modality=_lowerCamelCase				)





       def      _a       (   self				):
         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      _a       (   self				):
         UpperCamelCase_:   str        =							'resnet18'
         UpperCamelCase_:   Union[str, Any]        =							'microsoft/resnet-18'

         UpperCamelCase_:   List[Any]        =							AutoBackbone.from_pretrained(_lowerCamelCase    ,     use_timm_backbone=_lowerCamelCase				)
         UpperCamelCase_:   Tuple        =							AutoBackbone.from_pretrained(_lowerCamelCase				)

         self.assertEqual(len(timm_model.out_features				)    ,     len(transformers_model.out_features				)				)
         self.assertEqual(len(timm_model.stage_names				)    ,     len(transformers_model.stage_names				)				)
         self.assertEqual(timm_model.channels    ,     transformers_model.channels				)
         # Out indices are set to the last layer by default. For timm models, we don't know
         # the number of layers in advance, so we set it to (-1,), whereas for transformers
         # models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
         self.assertEqual(timm_model.out_indices    ,     (-1,)				)
         self.assertEqual(transformers_model.out_indices    ,     [len(timm_model.stage_names				) - 1]				)

         UpperCamelCase_:   Any        =							AutoBackbone.from_pretrained(_lowerCamelCase    ,     use_timm_backbone=_lowerCamelCase    ,     out_indices=[1, 2, 3]				)
         UpperCamelCase_:   Union[str, Any]        =							AutoBackbone.from_pretrained(_lowerCamelCase    ,     out_indices=[1, 2, 3]				)

         self.assertEqual(timm_model.out_indices    ,     transformers_model.out_indices				)
         self.assertEqual(len(timm_model.out_features				)    ,     len(transformers_model.out_features				)				)
         self.assertEqual(timm_model.channels    ,     transformers_model.channels				)





       @unittest.skip('TimmBackbone doesn\'t support feed forward chunking'				)
       def      _a       (   self				):
         pass





       @unittest.skip('TimmBackbone doesn\'t have num_hidden_layers attribute'				)
       def      _a       (   self				):
         pass





       @unittest.skip('TimmBackbone initialization is managed on the timm side'				)
       def      _a       (   self				):
         pass





       @unittest.skip('TimmBackbone models doesn\'t have inputs_embeds'				)
       def      _a       (   self				):
         pass





       @unittest.skip('TimmBackbone models doesn\'t have inputs_embeds'				)
       def      _a       (   self				):
         pass





       @unittest.skip('TimmBackbone model cannot be created without specifying a backbone checkpoint'				)
       def      _a       (   self				):
         pass





       @unittest.skip('Only checkpoints on timm can be loaded into TimmBackbone'				)
       def      _a       (   self				):
         pass





       @unittest.skip('model weights aren\'t tied in TimmBackbone.'				)
       def      _a       (   self				):
         pass





       @unittest.skip('model weights aren\'t tied in TimmBackbone.'				)
       def      _a       (   self				):
         pass





       @unittest.skip('Only checkpoints on timm can be loaded into TimmBackbone'				)
       def      _a       (   self				):
         pass





       @unittest.skip('Only checkpoints on timm can be loaded into TimmBackbone'				)
       def      _a       (   self				):
         pass





       @unittest.skip('TimmBackbone doesn\'t have hidden size info in its configuration.'				)
       def      _a       (   self				):
         pass





       @unittest.skip('TimmBackbone doesn\'t support output_attentions.'				)
       def      _a       (   self				):
         pass





       @unittest.skip('Safetensors is not supported by timm.'				)
       def      _a       (   self				):
         pass





       @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'				)
       def      _a       (   self				):
         pass





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   Tuple        =							self.model_tester.prepare_config_and_inputs_for_common()

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

           UpperCamelCase_:   int        =							['pixel_values']
           self.assertListEqual(arg_names[:1]    ,     _lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   Dict        =							self.model_tester.prepare_config_and_inputs_for_common()
         UpperCamelCase_:   Optional[Any]        =							True
         UpperCamelCase_:   Optional[Any]        =							self.has_attentions

         # no need to test all models as different heads yield the same functionality
         UpperCamelCase_:   Union[str, Any]        =							self.all_model_classes[0]
         UpperCamelCase_:   Tuple        =							model_class(_lowerCamelCase				)
         model.to(_lowerCamelCase				)

         UpperCamelCase_:   Union[str, Any]        =							self._prepare_for_class(_lowerCamelCase    ,     _lowerCamelCase				)
         UpperCamelCase_:   Tuple        =							model(**_lowerCamelCase				)
         UpperCamelCase_:   Union[str, Any]        =							outputs[0][-1]

         # Encoder-/Decoder-only models
         UpperCamelCase_:   List[str]        =							outputs.hidden_states[0]
         hidden_states.retain_grad()

         if self.has_attentions:
           UpperCamelCase_:   List[str]        =							outputs.attentions[0]
           attentions.retain_grad()

         output.flatten()[0].backward(retain_graph=_lowerCamelCase				)

         self.assertIsNotNone(hidden_states.grad				)

         if self.has_attentions:
           self.assertIsNotNone(attentions.grad				)





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   List[str]        =							self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           UpperCamelCase_:   List[str]        =							model_class(_lowerCamelCase				)
           model.to(_lowerCamelCase				)
           model.eval()
           UpperCamelCase_:   List[str]        =							model(**_lowerCamelCase				)

           self.assertEqual(len(result.feature_maps				)    ,     len(config.out_indices				)				)
           self.assertEqual(len(model.channels				)    ,     len(config.out_indices				)				)

           # Check output of last stage is taken if out_features=None, out_indices=None
           UpperCamelCase_:   Dict        =							copy.deepcopy(_lowerCamelCase				)
           UpperCamelCase_:   Union[str, Any]        =							None
           UpperCamelCase_:   Union[str, Any]        =							model_class(_lowerCamelCase				)
           model.to(_lowerCamelCase				)
           model.eval()
           UpperCamelCase_:   Dict        =							model(**_lowerCamelCase				)

           self.assertEqual(len(result.feature_maps				)    ,     1				)
           self.assertEqual(len(model.channels				)    ,     1				)

           # Check backbone can be initialized with fresh weights
           UpperCamelCase_:   Union[str, Any]        =							copy.deepcopy(_lowerCamelCase				)
           UpperCamelCase_:   List[Any]        =							False
           UpperCamelCase_:   Optional[Any]        =							model_class(_lowerCamelCase				)
           model.to(_lowerCamelCase				)
           model.eval()
           UpperCamelCase_:   List[Any]        =							model(**_lowerCamelCase				)







from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map
def 		_a       ( a					:Dict[str, torch.Tensor]					)				->     Dict[str, torch.Tensor]:
 a						    = []
 a						    = []
 a						    = []

 for rt in rc.restypes:
  a						    = rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
  restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names]					)
  a						    = {name: i for i, name in enumerate(a					)}
  restype_atomaa_to_atomaa_list.append(
      [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types]					)

  restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names]					)

 # Add dummy mapping for restype 'UNK'
 restype_atomaa_to_atomaa_list.append([0] * 14					)
 restype_atomaa_to_atomaa_list.append([0] * 37					)
 restype_atomaa_mask_list.append([0.0] * 14					)

 a						    = torch.tensor(
     a      ,				dtype=torch.intaa      ,				device=protein['''aatype'''].device      ,				)
 a						    = torch.tensor(
     a      ,				dtype=torch.intaa      ,				device=protein['''aatype'''].device      ,				)
 a						    = torch.tensor(
     a      ,				dtype=torch.floataa      ,				device=protein['''aatype'''].device      ,				)
 a						    = protein['''aatype'''].to(torch.long					)

 # create the mapping for (residx, atom14) --> atom37, i.e. an array
 # with shape (num_res, 14) containing the atom37 indices for this protein
 a						    = restype_atomaa_to_atomaa[protein_aatype]
 a						    = restype_atomaa_mask[protein_aatype]

 a						    = residx_atomaa_mask
 a						    = residx_atomaa_to_atomaa.long()

 # create the gather indices for mapping back
 a						    = restype_atomaa_to_atomaa[protein_aatype]
 a						    = residx_atomaa_to_atomaa.long()

 # create the corresponding mask
 a						    = torch.zeros([21, 37]      ,				dtype=torch.floataa      ,				device=protein['''aatype'''].device					)
 for restype, restype_letter in enumerate(rc.restypes					):
  a						    = rc.restype_atoa[restype_letter]
  a						    = rc.residue_atoms[restype_name]
  for atom_name in atom_names:
   a						    = rc.atom_order[atom_name]
   a						    = 1

 a						    = restype_atomaa_mask[protein_aatype]
 a						    = residx_atomaa_mask

 return protein

def 		_a       ( a					:Dict[str, torch.Tensor]					)				->     Dict[str, np.ndarray]:
 a						    = tree_map(lambda a					: torch.tensor(a      ,				device=batch['''aatype'''].device					)      ,				a      ,				np.ndarray					)
 a						    = tensor_tree_map(lambda a					: np.array(a					)      ,				make_atomaa_masks(a					)					)
 return out


