


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


							def __init__(self      ,				lowerCAmelCase      ,				lowerCAmelCase=1_3      ,				lowerCAmelCase=7      ,				lowerCAmelCase=False      ,				lowerCAmelCase=True      ,				lowerCAmelCase=False      ,				lowerCAmelCase=True      ,				lowerCAmelCase=3_3      ,				lowerCAmelCase=3_2      ,				lowerCAmelCase=5      ,				lowerCAmelCase=4      ,				lowerCAmelCase=3_7      ,				lowerCAmelCase="gelu"      ,				lowerCAmelCase=0.1      ,				lowerCAmelCase=0.1      ,				lowerCAmelCase=5_1_2      ,				lowerCAmelCase=1_6      ,				lowerCAmelCase=2      ,				lowerCAmelCase=0.02      ,				lowerCAmelCase=3      ,				lowerCAmelCase=4      ,				lowerCAmelCase=None      ,				):
									__lowercase=			parent
									__lowercase=			batch_size
									__lowercase=			seq_length
									__lowercase=			is_training
									__lowercase=			use_input_mask
									__lowercase=			use_token_type_ids
									__lowercase=			use_labels
									__lowercase=			vocab_size
									__lowercase=			hidden_size
									__lowercase=			num_hidden_layers
									__lowercase=			num_attention_heads
									__lowercase=			intermediate_size
									__lowercase=			hidden_act
									__lowercase=			hidden_dropout_prob
									__lowercase=			attention_probs_dropout_prob
									__lowercase=			max_position_embeddings
									__lowercase=			type_vocab_size
									__lowercase=			type_sequence_label_size
									__lowercase=			initializer_range
									__lowercase=			num_labels
									__lowercase=			num_choices
									__lowercase=			scope


							def 			_A       (self					):
									__lowercase=			ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size					)

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

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

									__lowercase=			self.get_config()

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


							def 			_A       (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 			_A       (self      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase					):
									__lowercase=			EsmModel(config=lowerCAmelCase					)
									model.to(lowerCAmelCase					)
									model.eval()
									__lowercase=			model(lowerCAmelCase      ,				attention_mask=lowerCAmelCase					)
									__lowercase=			model(lowerCAmelCase					)
									__lowercase=			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 			_A       (self      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase					):
									__lowercase=			EsmForMaskedLM(config=lowerCAmelCase					)
									model.to(lowerCAmelCase					)
									model.eval()
									__lowercase=			model(lowerCAmelCase      ,				attention_mask=lowerCAmelCase      ,				labels=lowerCAmelCase					)
									self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.vocab_size)					)


							def 			_A       (self      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase      ,				lowerCAmelCase					):
									__lowercase=			self.num_labels
									__lowercase=			EsmForTokenClassification(config=lowerCAmelCase					)
									model.to(lowerCAmelCase					)
									model.eval()
									__lowercase=			model(lowerCAmelCase      ,				attention_mask=lowerCAmelCase      ,				labels=lowerCAmelCase					)
									self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.num_labels)					)







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





@require_torch
class       A    (   A_					,				A_					,				unittest.TestCase ):
							UpperCamelCase_   :							List[Any]       =False

							UpperCamelCase_   :							int       =(
							    (
							        EsmForMaskedLM,
							        EsmModel,
							        EsmForSequenceClassification,
							        EsmForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							UpperCamelCase_   :							Dict       =()
							UpperCamelCase_   :							int       =(
							    {
							        '''feature-extraction''': EsmModel,
							        '''fill-mask''': EsmForMaskedLM,
							        '''text-classification''': EsmForSequenceClassification,
							        '''token-classification''': EsmForTokenClassification,
							        '''zero-shot''': EsmForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							UpperCamelCase_   :							Tuple       =True


							def 			_A       (self					):
									__lowercase=			EsmModelTester(self					)
									__lowercase=			ConfigTester(self      ,				config_class=lowerCAmelCase      ,				hidden_size=3_7					)


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


							def 			_A       (self					):
									__lowercase=			self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowerCAmelCase					)


							def 			_A       (self					):
									__lowercase=			self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
											__lowercase=			type
											self.model_tester.create_and_check_model(*lowerCAmelCase					)


							def 			_A       (self					):
									__lowercase=			self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase					)


							def 			_A       (self					):
									__lowercase=			self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase					)


							@slow
							def 			_A       (self					):
									for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__lowercase=			EsmModel.from_pretrained(lowerCAmelCase					)
											self.assertIsNotNone(lowerCAmelCase					)


							def 			_A       (self					):
									__lowercase=			self.model_tester.prepare_config_and_inputs()[0]
									__lowercase=			EsmEmbeddings(config=lowerCAmelCase					)

									__lowercase=			torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]					)
									__lowercase=			torch.as_tensor(
									    [
									        [
									            0 + model.padding_idx + 1,
									            1 + model.padding_idx + 1,
									            2 + model.padding_idx + 1,
									            model.padding_idx,
									        ]
									    ]					)
									__lowercase=			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 			_A       (self					):
									__lowercase=			self.model_tester.prepare_config_and_inputs()[0]
									__lowercase=			EsmEmbeddings(config=lowerCAmelCase					)

									__lowercase=			torch.empty(2      ,				4      ,				3_0					)
									__lowercase=			[
									    0 + embeddings.padding_idx + 1,
									    1 + embeddings.padding_idx + 1,
									    2 + embeddings.padding_idx + 1,
									    3 + embeddings.padding_idx + 1,
									]
									__lowercase=			torch.as_tensor([expected_single_positions, expected_single_positions]					)
									__lowercase=			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 			_A       (self					):
									pass


							@unittest.skip('Esm does not support embedding resizing'					)
							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





@require_torch
class       A    (   A_ ):


							@slow
							def 			_A       (self					):
									with torch.no_grad():
											__lowercase=			EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'					)
											model.eval()
											__lowercase=			torch.tensor([[0, 1, 2, 3, 4, 5]]					)
											__lowercase=			model(lowerCAmelCase					)[0]

											__lowercase=			3_3

											__lowercase=			torch.Size((1, 6, vocab_size)					)
											self.assertEqual(output.shape      ,				lowerCAmelCase					)

											__lowercase=			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]      ,				lowerCAmelCase      ,				atol=1E-4					)					)







							@slow
							def 			_A       (self					):
									with torch.no_grad():
											__lowercase=			EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'					)
											model.eval()

											__lowercase=			torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]					)
											__lowercase=			model(lowerCAmelCase					)[0]
											# compare the actual values for a slice.
											__lowercase=			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]      ,				lowerCAmelCase      ,				atol=1E-4					)					)










from __future__ import annotations







def        _lowerCamelCase(	lowercase__   , lowercase__	)						->	bool:

		'''simple docstring'''







		__lowercase=			get_failure_array(lowercase__	)

		# 2) Step through text searching for pattern
		__lowercase,							__lowercase=			0, 0  # index into text, pattern
		while i < len(lowercase__	):
				if pattern[j] == text[i]:
						if j == (len(lowercase__	) - 1):
								return True
						j += 1

				# if this is a prefix in our pattern
				# just go back far enough to continue
				elif j > 0:
						__lowercase=			failure[j - 1]
						continue
				i += 1
		return False







def        _lowerCamelCase(	lowercase__	)						->	list[int]:

		'''simple docstring'''







		__lowercase=			[0]
		__lowercase=			0
		__lowercase=			1
		while j < len(lowercase__	):
				if pattern[i] == pattern[j]:
						i += 1
				elif i > 0:
						__lowercase=			failure[i - 1]
						continue
				j += 1
				failure.append(lowercase__	)
		return failure


if __name__ == "__main__":
	# Test 1)
	lowerCAmelCase           =		'''abc1abc12'''
	lowerCAmelCase           =		'''alskfjaldsabc1abc1abc12k23adsfabcabc'''
	lowerCAmelCase           =		'''alskfjaldsk23adsfabcabc'''
	assert kmp(pattern, texta) and not kmp(pattern, texta)

	# Test 2)
	lowerCAmelCase           =		'''ABABX'''
	lowerCAmelCase           =		'''ABABZABABYABABX'''
	assert kmp(pattern, text)

	# Test 3)
	lowerCAmelCase           =		'''AAAB'''
	lowerCAmelCase           =		'''ABAAAAAB'''
	assert kmp(pattern, text)

	# Test 4)
	lowerCAmelCase           =		'''abcdabcy'''
	lowerCAmelCase           =		'''abcxabcdabxabcdabcdabcy'''
	assert kmp(pattern, text)

	# Test 5)
	lowerCAmelCase           =		'''aabaabaaa'''
	assert get_failure_array(pattern) == [0, 1, 0, 1, 2, 3, 4, 5, 2]






