

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        lowercase		:



 '''simple docstring'''


 def __init__(self			,			__a			,			__a=13			,			__a=7			,			__a=False			,			__a=True			,			__a=False			,			__a=True			,			__a=33			,			__a=32			,			__a=5			,			__a=4			,			__a=37			,			__a="gelu"			,			__a=0.1			,			__a=0.1			,			__a=512			,			__a=16			,			__a=2			,			__a=0.02			,			__a=3			,			__a=4			,			__a=None			,			)					->   str:






  """simple docstring"""





  UpperCAmelCase__     							=							parent
  UpperCAmelCase__     							=							batch_size
  UpperCAmelCase__     							=							seq_length
  UpperCAmelCase__     							=							is_training
  UpperCAmelCase__     							=							use_input_mask
  UpperCAmelCase__     							=							use_token_type_ids
  UpperCAmelCase__     							=							use_labels
  UpperCAmelCase__     							=							vocab_size
  UpperCAmelCase__     							=							hidden_size
  UpperCAmelCase__     							=							num_hidden_layers
  UpperCAmelCase__     							=							num_attention_heads
  UpperCAmelCase__     							=							intermediate_size
  UpperCAmelCase__     							=							hidden_act
  UpperCAmelCase__     							=							hidden_dropout_prob
  UpperCAmelCase__     							=							attention_probs_dropout_prob
  UpperCAmelCase__     							=							max_position_embeddings
  UpperCAmelCase__     							=							type_vocab_size
  UpperCAmelCase__     							=							type_sequence_label_size
  UpperCAmelCase__     							=							initializer_range
  UpperCAmelCase__     							=							num_labels
  UpperCAmelCase__     							=							num_choices
  UpperCAmelCase__     							=							scope




 def     UpperCamelCase__				(self					)					->   Optional[int]:






  """simple docstring"""





  UpperCAmelCase__     							=							ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size					)

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

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

  UpperCAmelCase__     							=							self.get_config()

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




 def     UpperCamelCase__				(self					)					->   int:






  """simple docstring"""





  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     UpperCamelCase__				(self			,			__a			,			__a			,			__a			,			__a			,			__a			,			__a					)					->   Optional[int]:






  """simple docstring"""





  UpperCAmelCase__     							=							EsmModel(config=_UpperCamelCase					)
  model.to(_UpperCamelCase					)
  model.eval()
  UpperCAmelCase__     							=							model(_UpperCamelCase			,			attention_mask=_UpperCamelCase					)
  UpperCAmelCase__     							=							model(_UpperCamelCase					)
  UpperCAmelCase__     							=							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     UpperCamelCase__				(self			,			__a			,			__a			,			__a			,			__a			,			__a			,			__a					)					->   Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__     							=							EsmForMaskedLM(config=_UpperCamelCase					)
  model.to(_UpperCamelCase					)
  model.eval()
  UpperCAmelCase__     							=							model(_UpperCamelCase			,			attention_mask=_UpperCamelCase			,			labels=_UpperCamelCase					)
  self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.vocab_size)					)




 def     UpperCamelCase__				(self			,			__a			,			__a			,			__a			,			__a			,			__a			,			__a					)					->   Dict:






  """simple docstring"""





  UpperCAmelCase__     							=							self.num_labels
  UpperCAmelCase__     							=							EsmForTokenClassification(config=_UpperCamelCase					)
  model.to(_UpperCamelCase					)
  model.eval()
  UpperCAmelCase__     							=							model(_UpperCamelCase			,			attention_mask=_UpperCamelCase			,			labels=_UpperCamelCase					)
  self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.num_labels)					)




 def     UpperCamelCase__				(self					)					->   Optional[int]:






  """simple docstring"""





  UpperCAmelCase__     							=							self.prepare_config_and_inputs()
  (
      (
      UpperCAmelCase__    
  )							,       (
      UpperCAmelCase__    
  )							,       (
      UpperCAmelCase__    
  )							,       (
      UpperCAmelCase__    
  )							,       (
      UpperCAmelCase__    
  )							,       (
      UpperCAmelCase__    
  )							,           
  )     							=							config_and_inputs
  UpperCAmelCase__     							=							{'input_ids': input_ids, 'attention_mask': input_mask}
  return config, inputs_dict







@require_torch
class        lowercase		(			snake_case__					,							snake_case__					,							unittest.TestCase   ):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE   				=							False

 __SCREAMING_SNAKE_CASE   				=							(
     (
         EsmForMaskedLM,
         EsmModel,
         EsmForSequenceClassification,
         EsmForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 __SCREAMING_SNAKE_CASE   				=							()
 __SCREAMING_SNAKE_CASE   				=							(
     {
         """feature-extraction""": EsmModel,
         """fill-mask""": EsmForMaskedLM,
         """text-classification""": EsmForSequenceClassification,
         """token-classification""": EsmForTokenClassification,
         """zero-shot""": EsmForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 __SCREAMING_SNAKE_CASE   				=							True




 def     UpperCamelCase__				(self					)					->   Any:






  """simple docstring"""





  UpperCAmelCase__     							=							EsmModelTester(self					)
  UpperCAmelCase__     							=							ConfigTester(self			,			config_class=_UpperCamelCase			,			hidden_size=37					)




 def     UpperCamelCase__				(self					)					->   List[Any]:






  """simple docstring"""





  self.config_tester.run_common_tests()




 def     UpperCamelCase__				(self					)					->   Optional[int]:






  """simple docstring"""





  UpperCAmelCase__     							=							self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*_UpperCamelCase					)




 def     UpperCamelCase__				(self					)					->   Tuple:






  """simple docstring"""





  UpperCAmelCase__     							=							self.model_tester.prepare_config_and_inputs()
  for type in ["absolute", "relative_key", "relative_key_query"]:
   UpperCAmelCase__     							=							type
   self.model_tester.create_and_check_model(*_UpperCamelCase					)




 def     UpperCamelCase__				(self					)					->   Tuple:






  """simple docstring"""





  UpperCAmelCase__     							=							self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase					)




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






  """simple docstring"""





  UpperCAmelCase__     							=							self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase					)




 @slow
 def     UpperCamelCase__				(self					)					->   Optional[Any]:






  """simple docstring"""





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




 def     UpperCamelCase__				(self					)					->   Tuple:






  """simple docstring"""





  UpperCAmelCase__     							=							self.model_tester.prepare_config_and_inputs()[0]
  UpperCAmelCase__     							=							EsmEmbeddings(config=_UpperCamelCase					)

  UpperCAmelCase__     							=							torch.as_tensor([[12, 31, 13, model.padding_idx]]					)
  UpperCAmelCase__     							=							torch.as_tensor(
      [
          [
              0 + model.padding_idx + 1,
              1 + model.padding_idx + 1,
              2 + model.padding_idx + 1,
              model.padding_idx,
          ]
      ]					)
  UpperCAmelCase__     							=							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     UpperCamelCase__				(self					)					->   Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__     							=							self.model_tester.prepare_config_and_inputs()[0]
  UpperCAmelCase__     							=							EsmEmbeddings(config=_UpperCamelCase					)

  UpperCAmelCase__     							=							torch.empty(2			,			4			,			30					)
  UpperCAmelCase__     							=							[
      0 + embeddings.padding_idx + 1,
      1 + embeddings.padding_idx + 1,
      2 + embeddings.padding_idx + 1,
      3 + embeddings.padding_idx + 1,
  ]
  UpperCAmelCase__     							=							torch.as_tensor([expected_single_positions, expected_single_positions]					)
  UpperCAmelCase__     							=							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     UpperCamelCase__				(self					)					->   Optional[Any]:






  """simple docstring"""





  pass




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






  """simple docstring"""





  pass




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






  """simple docstring"""





  pass







@require_torch
class        lowercase		(			snake_case__   ):



 '''simple docstring'''


 @slow
 def     UpperCamelCase__				(self					)					->   Union[str, Any]:






  """simple docstring"""





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

   UpperCAmelCase__     							=							33

   UpperCAmelCase__     							=							torch.Size((1, 6, vocab_size)					)
   self.assertEqual(output.shape			,			_UpperCamelCase					)

   UpperCAmelCase__     							=							torch.tensor(
       [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]					)
   self.assertTrue(torch.allclose(output[:, :3, :3]			,			_UpperCamelCase			,			atol=1E-4					)					)




 @slow
 def     UpperCamelCase__				(self					)					->   Optional[int]:






  """simple docstring"""





  with torch.no_grad():
   UpperCAmelCase__     							=							EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'					)
   model.eval()

   UpperCAmelCase__     							=							torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
   UpperCAmelCase__     							=							model(_UpperCamelCase					)[0]
   # compare the actual values for a slice.
   UpperCAmelCase__     							=							torch.tensor(
       [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]					)
   self.assertTrue(torch.allclose(output[:, :3, :3]			,			_UpperCamelCase			,			atol=1E-4					)					)






import enum
import os
from hashlib import shaaaa
from typing import Optional

from .. import config
from .logging import get_logger


UpperCamelCase__		 =     get_logger(__name__)

class     __snake_case    (				enum.Enum  ):
       """simple docstring"""







       __SCREAMING_SNAKE_CASE				    =							"all_checks"
       __SCREAMING_SNAKE_CASE				    =							"basic_checks"
       __SCREAMING_SNAKE_CASE				    =							"no_checks"






class     __snake_case    (				snake_case__  ):
       """simple docstring"""













class     __snake_case    (				snake_case__  ):
       """simple docstring"""













class     __snake_case    (				snake_case__  ):
       """simple docstring"""













class     __snake_case    (				snake_case__  ):
       """simple docstring"""














def 						lowerCamelCase__      ( __A    :Optional[dict]	,__A    :dict	,__A    :Any=None		):


       """simple docstring"""




       if expected_checksums is None:
              logger.info("""Unable to verify checksums."""		)
              return
       if len(set(__A		) - set(__A		)		) > 0:
              raise ExpectedMoreDownloadedFiles(str(set(__A		) - set(__A		)		)		)
       if len(set(__A		) - set(__A		)		) > 0:
              raise UnexpectedDownloadedFile(str(set(__A		) - set(__A		)		)		)
       __snake_case       	=							[url for url in expected_checksums if expected_checksums[url] != recorded_checksums[url]]
       __snake_case       	=							""" for """ + verification_name if verification_name is not None else """"""
       if len(__A		) > 0:
              raise NonMatchingChecksumError(
                  F'Checksums didn\'t match{for_verification_name}:\n'
                  F'{bad_urls}\n'
                  """Set `verification_mode='no_checks'` to skip checksums verification and ignore this error"""		)
       logger.info("""All the checksums matched successfully""" + for_verification_name		)

class     __snake_case    (				snake_case__  ):
       """simple docstring"""













class     __snake_case    (				snake_case__  ):
       """simple docstring"""













class     __snake_case    (				snake_case__  ):
       """simple docstring"""













class     __snake_case    (				snake_case__  ):
       """simple docstring"""














def 						lowerCamelCase__      ( __A    :Optional[dict]	,__A    :dict		):


       """simple docstring"""




       if expected_splits is None:
              logger.info("""Unable to verify splits sizes."""		)
              return
       if len(set(__A		) - set(__A		)		) > 0:
              raise ExpectedMoreSplits(str(set(__A		) - set(__A		)		)		)
       if len(set(__A		) - set(__A		)		) > 0:
              raise UnexpectedSplits(str(set(__A		) - set(__A		)		)		)
       __snake_case       	=							[
           {"""expected""": expected_splits[name], """recorded""": recorded_splits[name]}
           for name in expected_splits
           if expected_splits[name].num_examples != recorded_splits[name].num_examples
       ]
       if len(__A		) > 0:
              raise NonMatchingSplitsSizesError(str(__A		)		)
       logger.info("""All the splits matched successfully."""		)







def 						lowerCamelCase__      ( __A    :str	,__A    :bool = True		):


       """simple docstring"""




       if record_checksum:
              __snake_case       	=							shaaaa()
              with open(__A	,"""rb"""		) as f:
                     for chunk in iter(lambda: f.read(1 << 2_0		)	,b""""""		):
                            m.update(__A		)
                     __snake_case       	=							m.hexdigest()
       else:
              __snake_case       	=							None
       return {"num_bytes": os.path.getsize(__A		), "checksum": checksum}







def 						lowerCamelCase__      ( __A    :Optional[int]		):


       """simple docstring"""




       if dataset_size and config.IN_MEMORY_MAX_SIZE:
              return dataset_size < config.IN_MEMORY_MAX_SIZE
       else:
              return False

