




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






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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size )

							SCREAMING_SNAKE_CASE_   : Dict					  =      None
							if self.use_input_mask:
										SCREAMING_SNAKE_CASE_   : Tuple					  =      random_attention_mask([self.batch_size, self.seq_length] )

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

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.get_config()

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






				def 		snake_case	(			self ):
							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 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      EsmModel(config=snake_case__ )
							model.to(snake_case__ )
							model.eval()
							SCREAMING_SNAKE_CASE_   : Any					  =      model(snake_case__      ,attention_mask=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      model(snake_case__ )
							SCREAMING_SNAKE_CASE_   : Any					  =      model(snake_case__ )

							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 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Tuple					  =      EsmForMaskedLM(config=snake_case__ )
							model.to(snake_case__ )
							model.eval()
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      model(snake_case__      ,attention_mask=snake_case__      ,labels=snake_case__ )
							self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size) )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.num_labels
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      EsmForTokenClassification(config=snake_case__ )
							model.to(snake_case__ )
							model.eval()
							SCREAMING_SNAKE_CASE_   : Dict					  =      model(snake_case__      ,attention_mask=snake_case__      ,labels=snake_case__ )
							self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.num_labels) )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Any					  =      self.prepare_config_and_inputs()
							(
							    (
							    SCREAMING_SNAKE_CASE_
							)  ,    (
							    SCREAMING_SNAKE_CASE_
							)  ,    (
							    SCREAMING_SNAKE_CASE_
							)  ,    (
							    SCREAMING_SNAKE_CASE_
							)  ,    (
							    SCREAMING_SNAKE_CASE_
							)  ,    (
							    SCREAMING_SNAKE_CASE_
							)  ,    
							)   : Optional[int]					  =      config_and_inputs
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      {'input_ids': input_ids, 'attention_mask': input_mask}
							return config, inputs_dict







@require_torch
class 				lowerCAmelCase_       (					lowercase_						,     lowercase_						,     unittest.TestCase							):
				__a     :       Dict						=     False

				__a     :       Optional[Any]						=     (
				    (
				        EsmForMaskedLM,
				        EsmModel,
				        EsmForSequenceClassification,
				        EsmForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				__a     :       Optional[int]						=     ()
				__a     :       Any						=     (
				    {
				        """feature-extraction""": EsmModel,
				        """fill-mask""": EsmForMaskedLM,
				        """text-classification""": EsmForSequenceClassification,
				        """token-classification""": EsmForTokenClassification,
				        """zero-shot""": EsmForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__a     :       Any						=     True






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






				def 		snake_case	(			self ):
							self.config_tester.run_common_tests()






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
										SCREAMING_SNAKE_CASE_   : str					  =      type
										self.model_tester.create_and_check_model(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*snake_case__ )






				@slow
				def 		snake_case	(			self ):







							for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      EsmModel.from_pretrained(snake_case__ )
										self.assertIsNotNone(snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      self.model_tester.prepare_config_and_inputs()[0]
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      EsmEmbeddings(config=snake_case__ )

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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      self.model_tester.prepare_config_and_inputs()[0]
							SCREAMING_SNAKE_CASE_   : List[str]					  =      EsmEmbeddings(config=snake_case__ )

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






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






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






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







@require_torch
class 				lowerCAmelCase_       (					lowercase_							):






				@slow
				def 		snake_case	(			self ):







							with torch.no_grad():
										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D' )
										model.eval()
										SCREAMING_SNAKE_CASE_   : str					  =      torch.tensor([[0, 1, 2, 3, 4, 5]] )
										SCREAMING_SNAKE_CASE_   : List[str]					  =      model(snake_case__ )[0]

										SCREAMING_SNAKE_CASE_   : List[str]					  =      33

										SCREAMING_SNAKE_CASE_   : int					  =      torch.Size((1, 6, vocab_size) )
										self.assertEqual(output.shape      ,snake_case__ )

										SCREAMING_SNAKE_CASE_   : Dict					  =      torch.tensor(
										    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]] )
										self.assertTrue(torch.allclose(output[:, :3, :3]      ,snake_case__      ,atol=1E-4 ) )






				@slow
				def 		snake_case	(			self ):







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

										SCREAMING_SNAKE_CASE_   : List[Any]					  =      torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
										SCREAMING_SNAKE_CASE_   : Tuple					  =      model(snake_case__ )[0]
										# compare the actual values for a slice.
										SCREAMING_SNAKE_CASE_   : int					  =      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]      ,snake_case__      ,atol=1E-4 ) )







'''simple docstring'''

from __future__ import annotations

from collections import deque
class     __lowerCAmelCase		:





      '''simple docstring'''






      def __init__(self       :       str  ,		UpperCamelCase       :       list[str]					):

             '''simple docstring'''







             lowercase__      =					[]
             self.adlist.append(
                 {'''value''': '''''', '''next_states''': [], '''fail_state''': 0, '''output''': []}					)

             for keyword in keywords:
                    self.add_keyword(UpperCamelCase					)
             self.set_fail_transitions()



      def  UpperCamelCase__						(self       :       Optional[int]  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       str					):

             '''simple docstring'''







             for state in self.adlist[current_state]["next_states"]:
                    if char == self.adlist[state]["value"]:
                           return state
             return None



      def  UpperCamelCase__						(self       :       Optional[Any]  ,		UpperCamelCase       :       str					):

             '''simple docstring'''







             lowercase__      =					0
             for character in keyword:
                    lowercase__      =					self.find_next_state(UpperCamelCase  ,		UpperCamelCase					)
                    if next_state is None:
                           self.adlist.append(
                               {
                                   '''value''': character,
                                   '''next_states''': [],
                                   '''fail_state''': 0,
                                   '''output''': [],
                               }					)
                           self.adlist[current_state]["next_states"].append(len(self.adlist					) - 1					)
                           lowercase__      =					len(self.adlist					) - 1
                    else:
                           lowercase__      =					next_state
             self.adlist[current_state]["output"].append(UpperCamelCase					)



      def  UpperCamelCase__						(self       :       List[str]					):

             '''simple docstring'''







             lowercase__      =					deque()
             for node in self.adlist[0]["next_states"]:
                    q.append(UpperCamelCase					)
                    lowercase__      =					0
             while q:
                    lowercase__      =					q.popleft()
                    for child in self.adlist[r]["next_states"]:
                           q.append(UpperCamelCase					)
                           lowercase__      =					self.adlist[r]['''fail_state''']
                           while (
                               self.find_next_state(UpperCamelCase  ,		self.adlist[child]['''value''']					) is None
                               and state != 0
                           ):
                                  lowercase__      =					self.adlist[state]['''fail_state''']
                           lowercase__      =					self.find_next_state(
                               UpperCamelCase  ,		self.adlist[child]['''value''']					)
                           if self.adlist[child]["fail_state"] is None:
                                  lowercase__      =					0
                           lowercase__      =					(
                               self.adlist[child]['''output''']
                               + self.adlist[self.adlist[child]['''fail_state''']]['''output''']
                           )






      def  UpperCamelCase__						(self       :       Optional[int]  ,		UpperCamelCase       :       str					):

             '''simple docstring'''







             lowercase__      =					{}  # returns a dict with keywords and list of its occurrences
             lowercase__      =					0
             for i in range(len(UpperCamelCase					)					):
                    while (
                        self.find_next_state(UpperCamelCase  ,		string[i]					) is None
                        and current_state != 0
                    ):
                           lowercase__      =					self.adlist[current_state]['''fail_state''']
                    lowercase__      =					self.find_next_state(UpperCamelCase  ,		string[i]					)
                    if next_state is None:
                           lowercase__      =					0
                    else:
                           lowercase__      =					next_state
                           for key in self.adlist[current_state]["output"]:
                                  if key not in result:
                                         lowercase__      =					[]
                                  result[key].append(i - len(UpperCamelCase					) + 1					)
             return result


if __name__ == "__main__":
 import doctest

 doctest.testmod()




