
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.models.esm.modeling_esmfold import EsmForProteinFolding

class      _SCREAMING_SNAKE_CASE				:





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


								"""simple docstring"""

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





				def 							_UpperCAmelCase    (		self			:		List[Any]				):


								"""simple docstring"""

								A  :	Any				=					ids_tensor([self.batch_size, self.seq_length]       ,   self.vocab_size				)

								A  :	Optional[int]				=					None
								if self.use_input_mask:
												A  :	Optional[Any]				=					random_attention_mask([self.batch_size, self.seq_length]				)

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

								A  :	Any				=					self.get_config()

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





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


								"""simple docstring"""

								A  :	Any				=					EsmConfig(
								    vocab_size=33       ,   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       ,   is_folding_model=A_       ,   esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}       ,   )
								return config





				def 							_UpperCAmelCase    (		self			:		Optional[Any]       ,   snake_case_			:		str       ,   snake_case_			:		List[str]       ,   snake_case_			:		Optional[Any]       ,   snake_case_			:		int       ,   snake_case_			:		List[Any]       ,   snake_case_			:		int				):


								"""simple docstring"""

								A  :	Tuple				=					EsmForProteinFolding(config=A_				).float()
								model.to(A_				)
								model.eval()
								A  :	Tuple				=					model(A_       ,   attention_mask=A_				)
								A  :	List[Any]				=					model(A_				)
								A  :	Union[str, Any]				=					model(A_				)

								self.parent.assertEqual(result.positions.shape       ,   (8, self.batch_size, self.seq_length, 14, 3)				)
								self.parent.assertEqual(result.angles.shape       ,   (8, self.batch_size, self.seq_length, 7, 2)				)






				def 							_UpperCAmelCase    (		self			:		List[Any]				):


								"""simple docstring"""

								A  :	Any				=					self.prepare_config_and_inputs()
								(
								    (
								    A
								)			,							(
								    A
								)			,							(
								    A
								)			,							(
								    A
								)			,							(
								    A
								)			,							(
								    A
								)			,							
								)  :	str				=					config_and_inputs
								A  :	Tuple				=					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
								return config, inputs_dict



@require_torch
class      _SCREAMING_SNAKE_CASE				(		snake_case,     snake_case,     unittest.TestCase				):
				lowerCamelCase_	  =  False

				lowerCamelCase_	  =  (EsmForProteinFolding,) if is_torch_available() else ()
				lowerCamelCase_	  =  ()
				lowerCamelCase_	  =  {} if is_torch_available() else {}
				lowerCamelCase_	  =  False





				def 							_UpperCAmelCase    (		self			:		Any				):


								"""simple docstring"""

								A  :	List[str]				=					EsmFoldModelTester(self				)
								A  :	List[Any]				=					ConfigTester(self       ,   config_class=A_       ,   hidden_size=37				)





				def 							_UpperCAmelCase    (		self			:		Optional[int]				):


								"""simple docstring"""

								self.config_tester.run_common_tests()





				def 							_UpperCAmelCase    (		self			:		Optional[int]				):


								"""simple docstring"""

								A  :	Union[str, Any]				=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*A_				)





				@unittest.skip('''Does not support attention outputs'''				)
				def 							_UpperCAmelCase    (		self			:		Optional[Any]				):


								"""simple docstring"""

								pass





				@unittest.skip
				def 							_UpperCAmelCase    (		self			:		Tuple				):


								"""simple docstring"""

								pass





				@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			:		List[Any]				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold does not support passing input embeds!'''				)
				def 							_UpperCAmelCase    (		self			:		Dict				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold does not support head pruning.'''				)
				def 							_UpperCAmelCase    (		self			:		Optional[Any]				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold does not support head pruning.'''				)
				def 							_UpperCAmelCase    (		self			:		Tuple				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold does not support head pruning.'''				)
				def 							_UpperCAmelCase    (		self			:		Optional[Any]				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold does not support head pruning.'''				)
				def 							_UpperCAmelCase    (		self			:		Dict				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold does not support head pruning.'''				)
				def 							_UpperCAmelCase    (		self			:		Any				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold does not output hidden states in the normal way.'''				)
				def 							_UpperCAmelCase    (		self			:		List[str]				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMfold does not output hidden states in the normal way.'''				)
				def 							_UpperCAmelCase    (		self			:		str				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold only has one output format.'''				)
				def 							_UpperCAmelCase    (		self			:		Optional[Any]				):


								"""simple docstring"""

								pass





				@unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality'''				)
				def 							_UpperCAmelCase    (		self			:		str				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold does not support input chunking.'''				)
				def 							_UpperCAmelCase    (		self			:		str				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''				)
				def 							_UpperCAmelCase    (		self			:		List[str]				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''				)
				def 							_UpperCAmelCase    (		self			:		List[str]				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''				)
				def 							_UpperCAmelCase    (		self			:		Dict				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''				)
				def 							_UpperCAmelCase    (		self			:		Dict				):


								"""simple docstring"""

								pass





				@unittest.skip('''ESMFold doesn\'t support data parallel.'''				)
				def 							_UpperCAmelCase    (		self			:		List[Any]				):


								"""simple docstring"""

								pass






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


								"""simple docstring"""

								pass



@require_torch
class      _SCREAMING_SNAKE_CASE				(		snake_case				):





				@slow
				def 							_UpperCAmelCase    (		self			:		Optional[int]				):


								"""simple docstring"""

								A  :	int				=					EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''				).float()
								model.eval()
								A  :	str				=					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]				)
								A  :	int				=					model(A_				)['''positions''']
								A  :	Dict				=					torch.tensor([2.58_28, 0.79_93, -10.93_34]       ,   dtype=torch.floataa				)
								self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]       ,   A_       ,   atol=1E-4				)				)



"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import (
					    TFRoFormerForCausalLM,
					    TFRoFormerForMaskedLM,
					    TFRoFormerForMultipleChoice,
					    TFRoFormerForQuestionAnswering,
					    TFRoFormerForSequenceClassification,
					    TFRoFormerForTokenClassification,
					    TFRoFormerModel,
					)
					from transformers.models.roformer.modeling_tf_roformer import (
					    TFRoFormerSelfAttention,
					    TFRoFormerSinusoidalPositionalEmbedding,
					)




class 			__snake_case   :




							def __init__(     self:    Any     ,							A_:    str     ,							A_:    str=13     ,							A_:    Dict=7     ,							A_:    Optional[int]=True     ,							A_:    int=True     ,							A_:    List[str]=True     ,							A_:    Any=True     ,							A_:    List[Any]=99     ,							A_:    Tuple=32     ,							A_:    Union[str, Any]=2     ,							A_:    Dict=4     ,							A_:    Tuple=37     ,							A_:    Dict="gelu"     ,							A_:    List[Any]=0.1     ,							A_:    int=0.1     ,							A_:    Union[str, Any]=5_12     ,							A_:    List[Any]=16     ,							A_:    Optional[Any]=2     ,							A_:    int=0.02     ,							A_:    List[str]=3     ,							A_:    Optional[int]=4     ,							A_:    List[str]=None     ,							):
										__lowerCamelCase          =      parent
										__lowerCamelCase          =      13
										__lowerCamelCase          =      7
										__lowerCamelCase          =      True
										__lowerCamelCase          =      True
										__lowerCamelCase          =      True
										__lowerCamelCase          =      True
										__lowerCamelCase          =      99
										__lowerCamelCase          =      32
										__lowerCamelCase          =      2
										__lowerCamelCase          =      4
										__lowerCamelCase          =      37
										__lowerCamelCase          =      """gelu"""
										__lowerCamelCase          =      0.1
										__lowerCamelCase          =      0.1
										__lowerCamelCase          =      5_12
										__lowerCamelCase          =      16
										__lowerCamelCase          =      2
										__lowerCamelCase          =      0.02
										__lowerCamelCase          =      3
										__lowerCamelCase          =      4
										__lowerCamelCase          =      None




							def   __a   (     self:    Optional[Any]   ):
										__lowerCamelCase          =      ids_tensor([self.batch_size, self.seq_length]     ,							self.vocab_size   )

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

										__lowerCamelCase          =      None
										if self.use_token_type_ids:
													__lowerCamelCase          =      ids_tensor([self.batch_size, self.seq_length]     ,							self.type_vocab_size   )

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

										__lowerCamelCase          =      RoFormerConfig(
										    vocab_size=self.vocab_size     ,							hidden_size=self.hidden_size     ,							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     ,							return_dict=A_     ,							)

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




							def   __a   (     self:    Union[str, Any]     ,							A_:    str     ,							A_:    Dict     ,							A_:    Union[str, Any]     ,							A_:    Optional[int]     ,							A_:    Any     ,							A_:    Dict     ,							A_:    int   ):
										__lowerCamelCase          =      TFRoFormerModel(config=A_   )
										__lowerCamelCase          =      {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

										__lowerCamelCase          =      [input_ids, input_mask]
										__lowerCamelCase          =      model(A_   )

										__lowerCamelCase          =      model(A_   )

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




							def   __a   (     self:    List[str]     ,							A_:    int     ,							A_:    str     ,							A_:    Any     ,							A_:    int     ,							A_:    str     ,							A_:    str     ,							A_:    Union[str, Any]   ):
										__lowerCamelCase          =      True
										__lowerCamelCase          =      TFRoFormerForCausalLM(config=A_   )
										__lowerCamelCase          =      {
										    """input_ids""": input_ids,
										    """attention_mask""": input_mask,
										    """token_type_ids""": token_type_ids,
										}
										__lowerCamelCase          =      model(A_   )["""logits"""]
										self.parent.assertListEqual(
										    list(prediction_scores.numpy().shape   )     ,							[self.batch_size, self.seq_length, self.vocab_size]   )




							def   __a   (     self:    Any     ,							A_:    Optional[Any]     ,							A_:    Optional[Any]     ,							A_:    Union[str, Any]     ,							A_:    str     ,							A_:    str     ,							A_:    Optional[Any]     ,							A_:    int   ):
										__lowerCamelCase          =      TFRoFormerForMaskedLM(config=A_   )
										__lowerCamelCase          =      {
										    """input_ids""": input_ids,
										    """attention_mask""": input_mask,
										    """token_type_ids""": token_type_ids,
										}
										__lowerCamelCase          =      model(A_   )
										self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.seq_length, self.vocab_size)   )




							def   __a   (     self:    Optional[Any]     ,							A_:    Dict     ,							A_:    Union[str, Any]     ,							A_:    List[Any]     ,							A_:    str     ,							A_:    int     ,							A_:    Optional[Any]     ,							A_:    int   ):
										__lowerCamelCase          =      self.num_labels
										__lowerCamelCase          =      TFRoFormerForSequenceClassification(config=A_   )
										__lowerCamelCase          =      {
										    """input_ids""": input_ids,
										    """attention_mask""": input_mask,
										    """token_type_ids""": token_type_ids,
										}

										__lowerCamelCase          =      model(A_   )
										self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.num_labels)   )




							def   __a   (     self:    List[str]     ,							A_:    Tuple     ,							A_:    List[str]     ,							A_:    Union[str, Any]     ,							A_:    Optional[Any]     ,							A_:    List[str]     ,							A_:    Optional[int]     ,							A_:    List[Any]   ):
										__lowerCamelCase          =      self.num_choices
										__lowerCamelCase          =      TFRoFormerForMultipleChoice(config=A_   )
										__lowerCamelCase          =      tf.tile(tf.expand_dims(A_     ,							1   )     ,							(1, self.num_choices, 1)   )
										__lowerCamelCase          =      tf.tile(tf.expand_dims(A_     ,							1   )     ,							(1, self.num_choices, 1)   )
										__lowerCamelCase          =      tf.tile(tf.expand_dims(A_     ,							1   )     ,							(1, self.num_choices, 1)   )
										__lowerCamelCase          =      {
										    """input_ids""": multiple_choice_inputs_ids,
										    """attention_mask""": multiple_choice_input_mask,
										    """token_type_ids""": multiple_choice_token_type_ids,
										}
										__lowerCamelCase          =      model(A_   )
										self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.num_choices)   )




							def   __a   (     self:    Union[str, Any]     ,							A_:    Dict     ,							A_:    List[Any]     ,							A_:    List[Any]     ,							A_:    Union[str, Any]     ,							A_:    str     ,							A_:    Union[str, Any]     ,							A_:    List[str]   ):
										__lowerCamelCase          =      self.num_labels
										__lowerCamelCase          =      TFRoFormerForTokenClassification(config=A_   )
										__lowerCamelCase          =      {
										    """input_ids""": input_ids,
										    """attention_mask""": input_mask,
										    """token_type_ids""": token_type_ids,
										}
										__lowerCamelCase          =      model(A_   )
										self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.seq_length, self.num_labels)   )




							def   __a   (     self:    int     ,							A_:    Union[str, Any]     ,							A_:    Any     ,							A_:    Any     ,							A_:    Dict     ,							A_:    str     ,							A_:    int     ,							A_:    Tuple   ):
										__lowerCamelCase          =      TFRoFormerForQuestionAnswering(config=A_   )
										__lowerCamelCase          =      {
										    """input_ids""": input_ids,
										    """attention_mask""": input_mask,
										    """token_type_ids""": token_type_ids,
										}

										__lowerCamelCase          =      model(A_   )
										self.parent.assertEqual(result.start_logits.shape     ,							(self.batch_size, self.seq_length)   )
										self.parent.assertEqual(result.end_logits.shape     ,							(self.batch_size, self.seq_length)   )




							def   __a   (     self:    Union[str, Any]   ):
										__lowerCamelCase          =      self.prepare_config_and_inputs()
										(
										    (
										    __lowerCamelCase
										)     ,(
										    __lowerCamelCase
										)     ,(
										    __lowerCamelCase
										)     ,(
										    __lowerCamelCase
										)     ,(
										    __lowerCamelCase
										)     ,(
										    __lowerCamelCase
										)     ,(
										    __lowerCamelCase
										)     ,
										)          =      config_and_inputs
										__lowerCamelCase          =      {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
										return config, inputs_dict




@require_tf
class 			__snake_case   (lowerCamelCase			,		lowerCamelCase			,		unittest.TestCase  ):
							__a					 =       (
							    (
							        TFRoFormerModel,
							        TFRoFormerForCausalLM,
							        TFRoFormerForMaskedLM,
							        TFRoFormerForQuestionAnswering,
							        TFRoFormerForSequenceClassification,
							        TFRoFormerForTokenClassification,
							        TFRoFormerForMultipleChoice,
							    )
							    if is_tf_available()
							    else ()
							)
							__a					 =       (
							    {
							        '''feature-extraction''': TFRoFormerModel,
							        '''fill-mask''': TFRoFormerForMaskedLM,
							        '''question-answering''': TFRoFormerForQuestionAnswering,
							        '''text-classification''': TFRoFormerForSequenceClassification,
							        '''text-generation''': TFRoFormerForCausalLM,
							        '''token-classification''': TFRoFormerForTokenClassification,
							        '''zero-shot''': TFRoFormerForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)

							__a					 =       False
							__a					 =       False




							def   __a   (     self:    List[str]     ,							A_:    List[Any]     ,							A_:    List[str]     ,							A_:    str     ,							A_:    Union[str, Any]     ,							A_:    Dict   ):
										if pipeline_test_casse_name == "TextGenerationPipelineTests":
													return True

										return False




							def   __a   (     self:    Tuple   ):
										__lowerCamelCase          =      TFRoFormerModelTester(self   )
										__lowerCamelCase          =      ConfigTester(self     ,							config_class=A_     ,							hidden_size=37   )




							def   __a   (     self:    int   ):
										self.config_tester.run_common_tests()




							def   __a   (     self:    Optional[int]   ):
										__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*A_   )




							def   __a   (     self:    Union[str, Any]   ):
										__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*A_   )




							def   __a   (     self:    Optional[Any]   ):
										__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_lm_head(*A_   )




							def   __a   (     self:    Tuple   ):
										__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_multiple_choice(*A_   )




							def   __a   (     self:    Any   ):
										__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*A_   )




							def   __a   (     self:    int   ):
										__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*A_   )




							def   __a   (     self:    int   ):
										__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*A_   )




							@slow
							def   __a   (     self:    Optional[int]   ):
										__lowerCamelCase          =      TFRoFormerModel.from_pretrained("""junnyu/roformer_chinese_base"""   )
										self.assertIsNotNone(A_   )




@require_tf
class 			__snake_case   (unittest.TestCase  ):




							@slow
							def   __a   (     self:    Optional[Any]   ):
										__lowerCamelCase          =      TFRoFormerForMaskedLM.from_pretrained("""junnyu/roformer_chinese_base"""   )
										__lowerCamelCase          =      tf.constant([[0, 1, 2, 3, 4, 5]]   )
										__lowerCamelCase          =      model(A_   )[0]

										# TODO Replace vocab size
										__lowerCamelCase          =      5_00_00

										__lowerCamelCase          =      [1, 6, vocab_size]
										self.assertEqual(output.shape     ,							A_   )

										print(output[:, :3, :3]   )

										# TODO Replace values below with what was printed above.
										__lowerCamelCase          =      tf.constant(
										    [
										        [
										            [-0.12_053_341, -1.0_264_901, 0.29_221_946],
										            [-1.5_133_783, 0.197_433, 0.15_190_607],
										            [-5.0_135_403, -3.900_256, -0.84_038_764],
										        ]
										    ]   )
										tf.debugging.assert_near(output[:, :3, :3]     ,							A_     ,							atol=1E-4   )




@require_tf
class 			__snake_case   (unittest.TestCase  ):
							__a					 =       1e-4




							def   __a   (     self:    List[Any]   ):
										__lowerCamelCase          =      tf.constant([[4, 10]]   )
										__lowerCamelCase          =      TFRoFormerSinusoidalPositionalEmbedding(num_positions=6     ,							embedding_dim=6   )

										__lowerCamelCase          =      emba(input_ids.shape   )
										__lowerCamelCase          =      tf.constant(
										    [[0.0_000, 0.0_000, 0.0_000, 1.0_000, 1.0_000, 1.0_000], [0.8_415, 0.0_464, 0.0_022, 0.5_403, 0.9_989, 1.0_000]]   )

										tf.debugging.assert_near(A_     ,							A_     ,							atol=self.tolerance   )




							def   __a   (     self:    Optional[Any]   ):
										__lowerCamelCase          =      tf.constant(
										    [
										        [0.0_000, 0.0_000, 0.0_000, 0.0_000, 0.0_000],
										        [0.8_415, 0.8_219, 0.8_020, 0.7_819, 0.7_617],
										        [0.9_093, 0.9_364, 0.9_581, 0.9_749, 0.9_870],
										    ]   )
										__lowerCamelCase          =      TFRoFormerSinusoidalPositionalEmbedding(num_positions=5_12     ,							embedding_dim=5_12   )
										emba([2, 16, 5_12]   )
										__lowerCamelCase          =      emba.weight[:3, :5]

										tf.debugging.assert_near(A_     ,							A_     ,							atol=self.tolerance   )




@require_tf
class 			__snake_case   (unittest.TestCase  ):
							__a					 =       1e-4




							def   __a   (     self:    str   ):
										# 2,12,16,64
										__lowerCamelCase          =      tf.reshape(tf.range(2 * 12 * 16 * 64     ,							dtype=tf.floataa   )     ,							shape=(2, 12, 16, 64)   ) / 1_00

										__lowerCamelCase          =      -tf.reshape(tf.range(2 * 12 * 16 * 64     ,							dtype=tf.floataa   )     ,							shape=(2, 12, 16, 64)   ) / 1_00

										__lowerCamelCase          =      TFRoFormerSinusoidalPositionalEmbedding(num_positions=32     ,							embedding_dim=64   )
										__lowerCamelCase          =      embed_positions([2, 16, 7_68]   )[None, None, :, :]

										__lowerCamelCase     ,__lowerCamelCase          =      TFRoFormerSelfAttention.apply_rotary_position_embeddings(
										    A_     ,							A_     ,							A_   )

										__lowerCamelCase          =      tf.constant(
										    [
										        [0.0_000, 0.0_100, 0.0_200, 0.0_300, 0.0_400, 0.0_500, 0.0_600, 0.0_700],
										        [-0.2_012, 0.8_897, 0.0_263, 0.9_401, 0.2_074, 0.9_463, 0.3_481, 0.9_343],
										        [-1.7_057, 0.6_271, -1.2_145, 1.3_897, -0.6_303, 1.7_647, -0.1_173, 1.8_985],
										        [-2.1_731, -1.6_397, -2.7_358, 0.2_854, -2.1_840, 1.7_183, -1.3_018, 2.4_871],
										        [0.2_717, -3.6_173, -2.9_206, -2.1_988, -3.6_638, 0.3_858, -2.9_155, 2.2_980],
										        [3.9_859, -2.1_580, -0.7_984, -4.4_904, -4.1_181, -2.0_252, -4.4_782, 1.1_253],
										    ]   )
										__lowerCamelCase          =      tf.constant(
										    [
										        [0.0_000, -0.0_100, -0.0_200, -0.0_300, -0.0_400, -0.0_500, -0.0_600, -0.0_700],
										        [0.2_012, -0.8_897, -0.0_263, -0.9_401, -0.2_074, -0.9_463, -0.3_481, -0.9_343],
										        [1.7_057, -0.6_271, 1.2_145, -1.3_897, 0.6_303, -1.7_647, 0.1_173, -1.8_985],
										        [2.1_731, 1.6_397, 2.7_358, -0.2_854, 2.1_840, -1.7_183, 1.3_018, -2.4_871],
										        [-0.2_717, 3.6_173, 2.9_206, 2.1_988, 3.6_638, -0.3_858, 2.9_155, -2.2_980],
										        [-3.9_859, 2.1_580, 0.7_984, 4.4_904, 4.1_181, 2.0_252, 4.4_782, -1.1_253],
										    ]   )

										tf.debugging.assert_near(query_layer[0, 0, :6, :8]     ,							A_     ,							atol=self.tolerance   )
										tf.debugging.assert_near(key_layer[0, 0, :6, :8]     ,							A_     ,							atol=self.tolerance   )




