



'''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 				lowercase__      :





	def __init__(					self   :							List[Any]					,lowerCamelCase__   :							Tuple					,lowerCamelCase__   :							List[str]=13					,lowerCamelCase__   :							str=7					,lowerCamelCase__   :							Dict=False					,lowerCamelCase__   :							List[str]=True					,lowerCamelCase__   :							str=False					,lowerCamelCase__   :							Tuple=True					,lowerCamelCase__   :							Optional[int]=33					,lowerCamelCase__   :							Tuple=32					,lowerCamelCase__   :							Union[str, Any]=5					,lowerCamelCase__   :							List[Any]=4					,lowerCamelCase__   :							Tuple=37					,lowerCamelCase__   :							List[Any]="gelu"					,lowerCamelCase__   :							int=0.1					,lowerCamelCase__   :							List[Any]=0.1					,lowerCamelCase__   :							str=512					,lowerCamelCase__   :							str=16					,lowerCamelCase__   :							Optional[int]=2					,lowerCamelCase__   :							str=0.0_2					,lowerCamelCase__   :							Tuple=3					,lowerCamelCase__   :							int=4					,lowerCamelCase__   :							int=None					,):







							'''simple docstring'''
							_UpperCamelCase   :						List[str]												=  parent
							_UpperCamelCase   :						Union[str, Any]												=  batch_size
							_UpperCamelCase   :						Optional[Any]												=  seq_length
							_UpperCamelCase   :						Dict												=  is_training
							_UpperCamelCase   :						Tuple												=  use_input_mask
							_UpperCamelCase   :						List[str]												=  use_token_type_ids
							_UpperCamelCase   :						List[str]												=  use_labels
							_UpperCamelCase   :						Optional[Any]												=  vocab_size
							_UpperCamelCase   :						Dict												=  hidden_size
							_UpperCamelCase   :						Optional[int]												=  num_hidden_layers
							_UpperCamelCase   :						Optional[int]												=  num_attention_heads
							_UpperCamelCase   :						int												=  intermediate_size
							_UpperCamelCase   :						Any												=  hidden_act
							_UpperCamelCase   :						Dict												=  hidden_dropout_prob
							_UpperCamelCase   :						int												=  attention_probs_dropout_prob
							_UpperCamelCase   :						Optional[int]												=  max_position_embeddings
							_UpperCamelCase   :						str												=  type_vocab_size
							_UpperCamelCase   :						List[Any]												=  type_sequence_label_size
							_UpperCamelCase   :						Union[str, Any]												=  initializer_range
							_UpperCamelCase   :						int												=  num_labels
							_UpperCamelCase   :						int												=  num_choices
							_UpperCamelCase   :						Any												=  scope





	def       UpperCamelCase_    (					self   :							int					):







							'''simple docstring'''
							_UpperCamelCase   :						str												=  ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size					)

							_UpperCamelCase   :						Tuple												=  None
							if self.use_input_mask:
													_UpperCamelCase   :						Union[str, Any]												=  random_attention_mask([self.batch_size, self.seq_length]					)

							_UpperCamelCase   :						Any												=  None
							_UpperCamelCase   :						str												=  None
							_UpperCamelCase   :						List[str]												=  None
							if self.use_labels:
													_UpperCamelCase   :						Union[str, Any]												=  ids_tensor([self.batch_size]					,self.type_sequence_label_size					)
													_UpperCamelCase   :						Union[str, Any]												=  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       UpperCamelCase_    (					self   :							Union[str, Any]					):







							'''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   :							Any					,lowerCamelCase__   :							Optional[Any]					,lowerCamelCase__   :							Optional[int]					,lowerCamelCase__   :							Tuple					,lowerCamelCase__   :							Optional[Any]					,lowerCamelCase__   :							Any					,lowerCamelCase__   :							Tuple					):







							'''simple docstring'''
							_UpperCamelCase   :						int												=  EsmModel(config=lowerCamelCase__					)
							model.to(lowerCamelCase__					)
							model.eval()
							_UpperCamelCase   :						Tuple												=  model(lowerCamelCase__					,attention_mask=lowerCamelCase__					)
							_UpperCamelCase   :						Optional[int]												=  model(lowerCamelCase__					)
							_UpperCamelCase   :						str												=  model(lowerCamelCase__					)

							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   :							Optional[Any]					,lowerCamelCase__   :							List[str]					,lowerCamelCase__   :							Optional[int]					,lowerCamelCase__   :							Optional[Any]					,lowerCamelCase__   :							Dict					,lowerCamelCase__   :							Tuple					,lowerCamelCase__   :							Optional[Any]					):







							'''simple docstring'''
							_UpperCamelCase   :						Any												=  EsmForMaskedLM(config=lowerCamelCase__					)
							model.to(lowerCamelCase__					)
							model.eval()
							_UpperCamelCase   :						str												=  model(lowerCamelCase__					,attention_mask=lowerCamelCase__					,labels=lowerCamelCase__					)
							self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.vocab_size)					)





	def       UpperCamelCase_    (					self   :							Dict					,lowerCamelCase__   :							Union[str, Any]					,lowerCamelCase__   :							Tuple					,lowerCamelCase__   :							int					,lowerCamelCase__   :							List[str]					,lowerCamelCase__   :							List[str]					,lowerCamelCase__   :							Optional[Any]					):







							'''simple docstring'''
							_UpperCamelCase   :						Dict												=  self.num_labels
							_UpperCamelCase   :						Tuple												=  EsmForTokenClassification(config=lowerCamelCase__					)
							model.to(lowerCamelCase__					)
							model.eval()
							_UpperCamelCase   :						List[str]												=  model(lowerCamelCase__					,attention_mask=lowerCamelCase__					,labels=lowerCamelCase__					)
							self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.num_labels)					)





	def       UpperCamelCase_    (					self   :							List[str]					):







							'''simple docstring'''
							_UpperCamelCase   :						Optional[int]												=  self.prepare_config_and_inputs()
							(
							    (
							    _UpperCamelCase
							)							,     (
							    _UpperCamelCase
							)							,     (
							    _UpperCamelCase
							)							,     (
							    _UpperCamelCase
							)							,     (
							    _UpperCamelCase
							)							,     (
							    _UpperCamelCase
							)							,     
							)   :						int												=  config_and_inputs
							_UpperCamelCase   :						Dict												=  {'input_ids': input_ids, 'attention_mask': input_mask}
							return config, inputs_dict






@require_torch
class 				lowercase__      (				lowercase		,   lowercase		,   unittest.TestCase  ):
	lowercase__    							=     False

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





	def       UpperCamelCase_    (					self   :							List[Any]					):







							'''simple docstring'''
							_UpperCamelCase   :						Any												=  EsmModelTester(self					)
							_UpperCamelCase   :						Tuple												=  ConfigTester(self					,config_class=lowerCamelCase__					,hidden_size=37					)





	def       UpperCamelCase_    (					self   :							int					):







							'''simple docstring'''
							self.config_tester.run_common_tests()





	def       UpperCamelCase_    (					self   :							str					):







							'''simple docstring'''
							_UpperCamelCase   :						Optional[int]												=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowerCamelCase__					)





	def       UpperCamelCase_    (					self   :							Optional[Any]					):







							'''simple docstring'''
							_UpperCamelCase   :						List[Any]												=  self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
													_UpperCamelCase   :						Tuple												=  type
													self.model_tester.create_and_check_model(*lowerCamelCase__					)





	def       UpperCamelCase_    (					self   :							List[Any]					):







							'''simple docstring'''
							_UpperCamelCase   :						Union[str, Any]												=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase__					)





	def       UpperCamelCase_    (					self   :							Optional[int]					):







							'''simple docstring'''
							_UpperCamelCase   :						Optional[int]												=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__					)





	@slow
	def       UpperCamelCase_    (					self   :							Tuple					):







							'''simple docstring'''
							for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													_UpperCamelCase   :						List[str]												=  EsmModel.from_pretrained(lowerCamelCase__					)
													self.assertIsNotNone(lowerCamelCase__					)





	def       UpperCamelCase_    (					self   :							Dict					):







							'''simple docstring'''
							_UpperCamelCase   :						str												=  self.model_tester.prepare_config_and_inputs()[0]
							_UpperCamelCase   :						Union[str, Any]												=  EsmEmbeddings(config=lowerCamelCase__					)

							_UpperCamelCase   :						Dict												=  torch.as_tensor([[12, 31, 13, model.padding_idx]]					)
							_UpperCamelCase   :						List[str]												=  torch.as_tensor(
							    [
							        [
							            0 + model.padding_idx + 1,
							            1 + model.padding_idx + 1,
							            2 + model.padding_idx + 1,
							            model.padding_idx,
							        ]
							    ]					)
							_UpperCamelCase   :						Dict												=  create_position_ids_from_input_ids(lowerCamelCase__					,model.padding_idx					)
							self.assertEqual(position_ids.shape					,expected_positions.shape					)
							self.assertTrue(torch.all(torch.eq(lowerCamelCase__					,lowerCamelCase__					)					)					)





	def       UpperCamelCase_    (					self   :							List[Any]					):







							'''simple docstring'''
							_UpperCamelCase   :						List[str]												=  self.model_tester.prepare_config_and_inputs()[0]
							_UpperCamelCase   :						Optional[Any]												=  EsmEmbeddings(config=lowerCamelCase__					)

							_UpperCamelCase   :						Tuple												=  torch.empty(2					,4					,30					)
							_UpperCamelCase   :						Any												=  [
							    0 + embeddings.padding_idx + 1,
							    1 + embeddings.padding_idx + 1,
							    2 + embeddings.padding_idx + 1,
							    3 + embeddings.padding_idx + 1,
							]
							_UpperCamelCase   :						Dict												=  torch.as_tensor([expected_single_positions, expected_single_positions]					)
							_UpperCamelCase   :						int												=  embeddings.create_position_ids_from_inputs_embeds(lowerCamelCase__					)
							self.assertEqual(position_ids.shape					,expected_positions.shape					)
							self.assertTrue(torch.all(torch.eq(lowerCamelCase__					,lowerCamelCase__					)					)					)





	@unittest.skip('Esm does not support embedding resizing'					)
	def       UpperCamelCase_    (					self   :							int					):







							'''simple docstring'''
							pass





	@unittest.skip('Esm does not support embedding resizing'					)
	def       UpperCamelCase_    (					self   :							str					):







							'''simple docstring'''
							pass





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







							'''simple docstring'''
							pass






@require_torch
class 				lowercase__      (				lowercase  ):





	@slow
	def       UpperCamelCase_    (					self   :							int					):







							'''simple docstring'''
							with torch.no_grad():
													_UpperCamelCase   :						List[Any]												=  EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'					)
													model.eval()
													_UpperCamelCase   :						Any												=  torch.tensor([[0, 1, 2, 3, 4, 5]]					)
													_UpperCamelCase   :						Optional[Any]												=  model(lowerCamelCase__					)[0]

													_UpperCamelCase   :						Optional[Any]												=  33

													_UpperCamelCase   :						List[str]												=  torch.Size((1, 6, vocab_size)					)
													self.assertEqual(output.shape					,lowerCamelCase__					)

													_UpperCamelCase   :						str												=  torch.tensor(
													    [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]]					)
													self.assertTrue(torch.allclose(output[:, :3, :3]					,lowerCamelCase__					,atol=1E-4					)					)





	@slow
	def       UpperCamelCase_    (					self   :							Any					):







							'''simple docstring'''
							with torch.no_grad():
													_UpperCamelCase   :						Optional[Any]												=  EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'					)
													model.eval()

													_UpperCamelCase   :						str												=  torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
													_UpperCamelCase   :						str												=  model(lowerCamelCase__					)[0]
													# compare the actual values for a slice.
													_UpperCamelCase   :						Optional[Any]												=  torch.tensor(
													    [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]					)
													self.assertTrue(torch.allclose(output[:, :3, :3]					,lowerCamelCase__					,atol=1E-4					)					)







'''simple docstring'''
from typing import List

import numpy as np






def A__       (	UpperCAmelCase_				):
						_UpperCamelCase   :						Any												=  {key: len(UpperCAmelCase_				) for key, value in gen_kwargs.items() if isinstance(UpperCAmelCase_    ,    UpperCAmelCase_				)}
						if len(set(lists_lengths.values()				)				) > 1:
												raise RuntimeError(
												    (
												        'Sharding is ambiguous for this dataset: '
												        + 'we found several data sources lists of different lengths, and we don\'t know over which list we should parallelize:\n'
												        + '\n'.join(f'\t- key {key} has length {length}' for key, length in lists_lengths.items()				)
												        + '\nTo fix this, check the \'gen_kwargs\' and make sure to use lists only for data sources, '
												        + 'and use tuples otherwise. In the end there should only be one single list, or several lists with the same length.'
												    )				)
						_UpperCamelCase   :						Tuple												=  max(lists_lengths.values()    ,    default=0				)
						return max(1    ,    UpperCAmelCase_				)






def A__       (	UpperCAmelCase_    ,    UpperCAmelCase_				):
						_UpperCamelCase   :						Dict												=  []
						for group_idx in range(UpperCAmelCase_				):
												_UpperCamelCase   :						int												=  num_shards // max_num_jobs + (group_idx < (num_shards % max_num_jobs))
												if num_shards_to_add == 0:
																		break
												_UpperCamelCase   :						Tuple												=  shards_indices_per_group[-1].stop if shards_indices_per_group else 0
												_UpperCamelCase   :						Tuple												=  range(UpperCAmelCase_    ,    start + num_shards_to_add				)
												shards_indices_per_group.append(UpperCAmelCase_				)
						return shards_indices_per_group






def A__       (	UpperCAmelCase_    ,    UpperCAmelCase_				):
						_UpperCamelCase   :						Tuple												=  _number_of_shards_in_gen_kwargs(UpperCAmelCase_				)
						if num_shards == 1:
												return [dict(UpperCAmelCase_				)]
						else:
												_UpperCamelCase   :						str												=  _distribute_shards(num_shards=UpperCAmelCase_    ,    max_num_jobs=UpperCAmelCase_				)
												return [
												    {
												        key: [value[shard_idx] for shard_idx in shard_indices_per_group[group_idx]]
												        if isinstance(UpperCAmelCase_    ,    UpperCAmelCase_				)
												        else value
												        for key, value in gen_kwargs.items()
												    }
												    for group_idx in range(len(UpperCAmelCase_				)				)
												]






def A__       (	UpperCAmelCase_				):
						return {
						    key: [value for gen_kwargs in gen_kwargs_list for value in gen_kwargs[key]]
						    if isinstance(gen_kwargs_list[0][key]    ,    UpperCAmelCase_				)
						    else gen_kwargs_list[0][key]
						    for key in gen_kwargs_list[0]
						}






def A__       (	UpperCAmelCase_    ,    UpperCAmelCase_				):
						_UpperCamelCase   :						int												=  {len(UpperCAmelCase_				) for value in gen_kwargs.values() if isinstance(UpperCAmelCase_    ,    UpperCAmelCase_				)}
						_UpperCamelCase   :						Union[str, Any]												=  {}
						for size in list_sizes:
												_UpperCamelCase   :						str												=  list(range(UpperCAmelCase_				)				)
												rng.shuffle(indices_per_size[size]				)
						# Now let's copy the gen_kwargs and shuffle the lists based on their sizes
						_UpperCamelCase   :						Union[str, Any]												=  dict(UpperCAmelCase_				)
						for key, value in shuffled_kwargs.items():
												if isinstance(UpperCAmelCase_    ,    UpperCAmelCase_				):
																		_UpperCamelCase   :						Dict												=  [value[i] for i in indices_per_size[len(UpperCAmelCase_				)]]
						return shuffled_kwargs


