


"""simple docstring"""




import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
   from transformers.models.esm.modeling_esm import (
       ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
       EsmEmbeddings,
       create_position_ids_from_input_ids,
   )





class __snake_case		:







   """simple docstring"""




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




   def SCREAMING_SNAKE_CASE_   (		self				:Tuple							):
          _a							 =					ids_tensor([self.batch_size, self.seq_length]					,	self.vocab_size							)

          _a							 =					None
          if self.use_input_mask:
                 _a							 =					random_attention_mask([self.batch_size, self.seq_length]							)

          _a							 =					None
          _a							 =					None
          _a							 =					None
          if self.use_labels:
                 _a							 =					ids_tensor([self.batch_size]					,	self.type_sequence_label_size							)
                 _a							 =					ids_tensor([self.batch_size, self.seq_length]					,	self.num_labels							)
                 _a							 =					ids_tensor([self.batch_size]					,	self.num_choices							)

          _a							 =					self.get_config()

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




   def SCREAMING_SNAKE_CASE_   (		self				:Optional[Any]							):
          return EsmConfig(
              vocab_size=self.vocab_size					,	hidden_size=self.hidden_size					,	pad_token_id=1					,	num_hidden_layers=self.num_hidden_layers					,	num_attention_heads=self.num_attention_heads					,	intermediate_size=self.intermediate_size					,	hidden_act=self.hidden_act					,	hidden_dropout_prob=self.hidden_dropout_prob					,	attention_probs_dropout_prob=self.attention_probs_dropout_prob					,	max_position_embeddings=self.max_position_embeddings					,	type_vocab_size=self.type_vocab_size					,	initializer_range=self.initializer_range					,	)




   def SCREAMING_SNAKE_CASE_   (		self				:Union[str, Any]					,	UpperCamelCase__				:Union[str, Any]					,	UpperCamelCase__				:Union[str, Any]					,	UpperCamelCase__				:Dict					,	UpperCamelCase__				:Optional[Any]					,	UpperCamelCase__				:int					,	UpperCamelCase__				:Tuple							):
          _a							 =					EsmModel(config=UpperCamelCase__							)
          model.to(UpperCamelCase__							)
          model.eval()
          _a							 =					model(UpperCamelCase__					,	attention_mask=UpperCamelCase__							)
          _a							 =					model(UpperCamelCase__							)
          _a							 =					model(UpperCamelCase__							)

          self.parent.assertEqual(result.last_hidden_state.shape					,	(self.batch_size, self.seq_length, self.hidden_size)							)
          self.parent.assertEqual(result.pooler_output.shape					,	(self.batch_size, self.hidden_size)							)




   def SCREAMING_SNAKE_CASE_   (		self				:Union[str, Any]					,	UpperCamelCase__				:Any					,	UpperCamelCase__				:Optional[Any]					,	UpperCamelCase__				:Union[str, Any]					,	UpperCamelCase__				:Union[str, Any]					,	UpperCamelCase__				:str					,	UpperCamelCase__				:int							):
          _a							 =					EsmForMaskedLM(config=UpperCamelCase__							)
          model.to(UpperCamelCase__							)
          model.eval()
          _a							 =					model(UpperCamelCase__					,	attention_mask=UpperCamelCase__					,	labels=UpperCamelCase__							)
          self.parent.assertEqual(result.logits.shape					,	(self.batch_size, self.seq_length, self.vocab_size)							)




   def SCREAMING_SNAKE_CASE_   (		self				:Optional[Any]					,	UpperCamelCase__				:Union[str, Any]					,	UpperCamelCase__				:Tuple					,	UpperCamelCase__				:int					,	UpperCamelCase__				:Optional[Any]					,	UpperCamelCase__				:Union[str, Any]					,	UpperCamelCase__				:str							):
          _a							 =					self.num_labels
          _a							 =					EsmForTokenClassification(config=UpperCamelCase__							)
          model.to(UpperCamelCase__							)
          model.eval()
          _a							 =					model(UpperCamelCase__					,	attention_mask=UpperCamelCase__					,	labels=UpperCamelCase__							)
          self.parent.assertEqual(result.logits.shape					,	(self.batch_size, self.seq_length, self.num_labels)							)




   def SCREAMING_SNAKE_CASE_   (		self				:Dict							):
          _a							 =					self.prepare_config_and_inputs()
          (
              (
              _a       
          ) ,			(
              _a       
          ) ,			(
              _a       
          ) ,			(
              _a       
          ) ,			(
              _a       
          ) ,			(
              _a       
          ) ,			       
          )							 =					config_and_inputs
          _a							 =					{"input_ids": input_ids, "attention_mask": input_mask}
          return config, inputs_dict






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







   """simple docstring"""




   lowerCAmelCase_					: Tuple  						=      False

   lowerCAmelCase_					: Any  						=      (
       (
           EsmForMaskedLM,
           EsmModel,
           EsmForSequenceClassification,
           EsmForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   lowerCAmelCase_					: List[str]  						=      ()
   lowerCAmelCase_					: Union[str, Any]  						=      (
       {
           '''feature-extraction''': EsmModel,
           '''fill-mask''': EsmForMaskedLM,
           '''text-classification''': EsmForSequenceClassification,
           '''token-classification''': EsmForTokenClassification,
           '''zero-shot''': EsmForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   lowerCAmelCase_					: Tuple  						=      True




   def SCREAMING_SNAKE_CASE_   (		self				:Optional[Any]							):
          _a							 =					EsmModelTester(self							)
          _a							 =					ConfigTester(self					,	config_class=UpperCamelCase__					,	hidden_size=37							)




   def SCREAMING_SNAKE_CASE_   (		self				:Dict							):
          self.config_tester.run_common_tests()




   def SCREAMING_SNAKE_CASE_   (		self				:List[Any]							):
          _a							 =					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*UpperCamelCase__							)




   def SCREAMING_SNAKE_CASE_   (		self				:Dict							):
          _a							 =					self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
                 _a							 =					type
                 self.model_tester.create_and_check_model(*UpperCamelCase__							)




   def SCREAMING_SNAKE_CASE_   (		self				:Optional[Any]							):
          _a							 =					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__							)




   def SCREAMING_SNAKE_CASE_   (		self				:List[str]							):
          _a							 =					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__							)




   @slow
   def SCREAMING_SNAKE_CASE_   (		self				:List[str]							):

          for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _a							 =					EsmModel.from_pretrained(UpperCamelCase__							)
                 self.assertIsNotNone(UpperCamelCase__							)




   def SCREAMING_SNAKE_CASE_   (		self				:Dict							):
          _a							 =					self.model_tester.prepare_config_and_inputs()[0]
          _a							 =					EsmEmbeddings(config=UpperCamelCase__							)

          _a							 =					torch.as_tensor([[12, 31, 13, model.padding_idx]]							)
          _a							 =					torch.as_tensor(
              [
                  [
                      0 + model.padding_idx + 1,
                      1 + model.padding_idx + 1,
                      2 + model.padding_idx + 1,
                      model.padding_idx,
                  ]
              ]							)
          _a							 =					create_position_ids_from_input_ids(UpperCamelCase__					,	model.padding_idx							)
          self.assertEqual(position_ids.shape					,	expected_positions.shape							)
          self.assertTrue(torch.all(torch.eq(UpperCamelCase__					,	UpperCamelCase__							)							)							)




   def SCREAMING_SNAKE_CASE_   (		self				:int							):
          _a							 =					self.model_tester.prepare_config_and_inputs()[0]
          _a							 =					EsmEmbeddings(config=UpperCamelCase__							)

          _a							 =					torch.empty(2					,	4					,	30							)
          _a							 =					[
              0 + embeddings.padding_idx + 1,
              1 + embeddings.padding_idx + 1,
              2 + embeddings.padding_idx + 1,
              3 + embeddings.padding_idx + 1,
          ]
          _a							 =					torch.as_tensor([expected_single_positions, expected_single_positions]							)
          _a							 =					embeddings.create_position_ids_from_inputs_embeds(UpperCamelCase__							)
          self.assertEqual(position_ids.shape					,	expected_positions.shape							)
          self.assertTrue(torch.all(torch.eq(UpperCamelCase__					,	UpperCamelCase__							)							)							)




   @unittest.skip("Esm does not support embedding resizing"							)
   def SCREAMING_SNAKE_CASE_   (		self				:List[Any]							):
          pass




   @unittest.skip("Esm does not support embedding resizing"							)
   def SCREAMING_SNAKE_CASE_   (		self				:Optional[int]							):
          pass




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






@require_torch
class __snake_case		(       _a   ):







   """simple docstring"""




   @slow
   def SCREAMING_SNAKE_CASE_   (		self				:int							):

          with torch.no_grad():
                 _a							 =					EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"							)
                 model.eval()
                 _a							 =					torch.tensor([[0, 1, 2, 3, 4, 5]]							)
                 _a							 =					model(UpperCamelCase__							)[0]

                 _a							 =					33

                 _a							 =					torch.Size((1, 6, vocab_size)							)
                 self.assertEqual(output.shape					,	UpperCamelCase__							)

                 _a							 =					torch.tensor(
                     [[[8.9215, -10.5_898, -6.4671], [-6.3967, -13.9_114, -1.1212], [-7.7812, -13.9_516, -3.7406]]]							)
                 self.assertTrue(torch.allclose(output[:, :3, :3]					,	UpperCamelCase__					,	atol=1E-4							)							)




   @slow
   def SCREAMING_SNAKE_CASE_   (		self				:int							):

          with torch.no_grad():
                 _a							 =					EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"							)
                 model.eval()

                 _a							 =					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
                 _a							 =					model(UpperCamelCase__							)[0]
                 # compare the actual values for a slice.
                 _a							 =					torch.tensor(
                     [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]							)
                 self.assertTrue(torch.allclose(output[:, :3, :3]					,	UpperCamelCase__					,	atol=1E-4							)							)








class 					SCREAMING_SNAKE_CASE_  :
       """simple docstring"""


       def __init__(				self	:List[Any]):



          """simple docstring"""

          _lowercase              ={}





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          print(self.vertex)
          for i in self.vertex:
             print(snake_case,      ' -> ',      ' -> '.join([str(snake_case) for j in self.vertex[i]]))





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:int,      snake_case	:int):



          """simple docstring"""

          if from_vertex in self.vertex:
             self.vertex[from_vertex].append(snake_case)
          else:
             # else make a new vertex
             _lowercase              =[to_vertex]





       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          _lowercase              =[False] * len(self.vertex)

          # call the recursive helper function
          for i in range(len(self.vertex)):
             if not visited[i]:
                self.dfs_recursive(snake_case,      snake_case)





       def 				UpperCamelCase__						(				self	:List[str],      snake_case	:int,      snake_case	:list):



          """simple docstring"""

          _lowercase              =True

          print(snake_case,      end=' ')

          # Recur for all the vertices that are adjacent to this node
          for i in self.vertex:
             if not visited[i]:
                self.dfs_recursive(snake_case,      snake_case)


if __name__ == "__main__":
    _SCREAMING_SNAKE_CASE   =  Graph()
    g.add_edge(0, 1)
    g.add_edge(0, 2)
    g.add_edge(1, 2)
    g.add_edge(2, 0)
    g.add_edge(2, 3)
    g.add_edge(3, 3)

    g.print_graph()
    print("DFS:")
    g.dfs()

    # OUTPUT:
    # 0  ->  1 -> 2
    # 1  ->  2
    # 2  ->  0 -> 3
    # 3  ->  3
    # DFS:
    #  0 1 2 3


