




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






						def __init__(     self	,    _A	,    _A=1_3	,    _A=7	,    _A=True	,    _A=True	,    _A=True	,    _A=True	,    _A=9_9	,    _A=3_2	,    _A=2	,    _A=4	,    _A=3_7	,    _A="gelu"	,    _A=0.1	,    _A=0.1	,    _A=5_1_2	,    _A=1_6	,    _A=2	,    _A=0.02	,    _A=3	,    _A=4	,    _A=None	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = parent
											__lowerCAmelCase            = 1_3
											__lowerCAmelCase            = 7
											__lowerCAmelCase            = True
											__lowerCAmelCase            = True
											__lowerCAmelCase            = True
											__lowerCAmelCase            = True
											__lowerCAmelCase            = 9_9
											__lowerCAmelCase            = 3_2
											__lowerCAmelCase            = 2
											__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






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__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        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__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        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__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        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__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        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__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        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__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        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__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        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__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        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__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     a__					(      snake_case__				,       snake_case__				,       unittest.TestCase   ):
						_a							:					Union[str, Any]      =					(
						    (
						        TFRoFormerModel,
						        TFRoFormerForCausalLM,
						        TFRoFormerForMaskedLM,
						        TFRoFormerForQuestionAnswering,
						        TFRoFormerForSequenceClassification,
						        TFRoFormerForTokenClassification,
						        TFRoFormerForMultipleChoice,
						    )
						    if is_tf_available()
						    else ()
						)
						_a							:					Dict      =					(
						    {
						        """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							:					Optional[Any]      =					False
						_a							:					List[str]      =					False






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											if pipeline_test_casse_name == "TextGenerationPipelineTests":
																return True

											return False






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFRoFormerModelTester(self    )
											__lowerCAmelCase            = ConfigTester(self	,    config_class=_A	,    hidden_size=3_7    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											self.config_tester.run_common_tests()






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_lm_head(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_multiple_choice(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*_A    )






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFRoFormerModel.from_pretrained("junnyu/roformer_chinese_base"    )
											self.assertIsNotNone(_A    )






@require_tf
class     a__					(      unittest.TestCase   ):






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__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_0_0_0_0

											__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_05_33_41, -1.0_26_49_01, 0.29_22_19_46],
											            [-1.5_13_37_83, 0.19_74_33, 0.15_19_06_07],
											            [-5.0_13_54_03, -3.90_02_56, -0.84_03_87_64],
											        ]
											    ]    )
											tf.debugging.assert_near(output[:, :3, :3]	,    _A	,    atol=1E-4    )






@require_tf
class     a__					(      unittest.TestCase   ):
						_a							:					str      =					1E-4






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = tf.constant([[4, 1_0]]    )
											__lowerCAmelCase            = TFRoFormerSinusoidalPositionalEmbedding(num_positions=6	,    embedding_dim=6    )

											__lowerCAmelCase            = emba(input_ids.shape    )
											__lowerCAmelCase            = tf.constant(
											    [[0.00_00, 0.00_00, 0.00_00, 1.00_00, 1.00_00, 1.00_00], [0.84_15, 0.04_64, 0.00_22, 0.54_03, 0.99_89, 1.00_00]]    )

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






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = tf.constant(
											    [
											        [0.00_00, 0.00_00, 0.00_00, 0.00_00, 0.00_00],
											        [0.84_15, 0.82_19, 0.80_20, 0.78_19, 0.76_17],
											        [0.90_93, 0.93_64, 0.95_81, 0.97_49, 0.98_70],
											    ]    )
											__lowerCAmelCase            = TFRoFormerSinusoidalPositionalEmbedding(num_positions=5_1_2	,    embedding_dim=5_1_2    )
											emba([2, 1_6, 5_1_2]    )
											__lowerCAmelCase            = emba.weight[:3, :5]

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






@require_tf
class     a__					(      unittest.TestCase   ):
						_a							:					Any      =					1E-4






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = tf.reshape(tf.range(2 * 1_2 * 1_6 * 6_4	,    dtype=tf.floataa    )	,    shape=(2, 1_2, 1_6, 6_4)    ) / 1_0_0

											__lowerCAmelCase            = -tf.reshape(tf.range(2 * 1_2 * 1_6 * 6_4	,    dtype=tf.floataa    )	,    shape=(2, 1_2, 1_6, 6_4)    ) / 1_0_0

											__lowerCAmelCase            = TFRoFormerSinusoidalPositionalEmbedding(num_positions=3_2	,    embedding_dim=6_4    )
											__lowerCAmelCase            = embed_positions([2, 1_6, 7_6_8]    )[None, None, :, :]

											__lowerCAmelCase     ,							__lowerCAmelCase            = TFRoFormerSelfAttention.apply_rotary_position_embeddings(
											    _A	,    _A	,    _A    )

											__lowerCAmelCase            = tf.constant(
											    [
											        [0.00_00, 0.01_00, 0.02_00, 0.03_00, 0.04_00, 0.05_00, 0.06_00, 0.07_00],
											        [-0.20_12, 0.88_97, 0.02_63, 0.94_01, 0.20_74, 0.94_63, 0.34_81, 0.93_43],
											        [-1.70_57, 0.62_71, -1.21_45, 1.38_97, -0.63_03, 1.76_47, -0.11_73, 1.89_85],
											        [-2.17_31, -1.63_97, -2.73_58, 0.28_54, -2.18_40, 1.71_83, -1.30_18, 2.48_71],
											        [0.27_17, -3.61_73, -2.92_06, -2.19_88, -3.66_38, 0.38_58, -2.91_55, 2.29_80],
											        [3.98_59, -2.15_80, -0.79_84, -4.49_04, -4.11_81, -2.02_52, -4.47_82, 1.12_53],
											    ]    )
											__lowerCAmelCase            = tf.constant(
											    [
											        [0.00_00, -0.01_00, -0.02_00, -0.03_00, -0.04_00, -0.05_00, -0.06_00, -0.07_00],
											        [0.20_12, -0.88_97, -0.02_63, -0.94_01, -0.20_74, -0.94_63, -0.34_81, -0.93_43],
											        [1.70_57, -0.62_71, 1.21_45, -1.38_97, 0.63_03, -1.76_47, 0.11_73, -1.89_85],
											        [2.17_31, 1.63_97, 2.73_58, -0.28_54, 2.18_40, -1.71_83, 1.30_18, -2.48_71],
											        [-0.27_17, 3.61_73, 2.92_06, 2.19_88, 3.66_38, -0.38_58, 2.91_55, -2.29_80],
											        [-3.98_59, 2.15_80, 0.79_84, 4.49_04, 4.11_81, 2.02_52, 4.47_82, -1.12_53],
											    ]    )

											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    )










from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__					    =    logging.get_logger(__name__)

UpperCamelCase__					    =    {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class     a__					(      snake_case__   ):
						_a							:					Tuple      =					"""nllb-moe"""
						_a							:					Dict      =					["""past_key_values"""]
						_a							:					Optional[int]      =					{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






						def __init__(     self	,    _A=1_2_8_1_1_2	,    _A=1_0_2_4	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=1_2	,    _A=4_0_9_6	,    _A=1_6	,    _A=0.05	,    _A=0.05	,    _A=True	,    _A=True	,    _A="relu"	,    _A=1_0_2_4	,    _A=0.1	,    _A=0.1	,    _A=0.0	,    _A=0.02	,    _A=2	,    _A=True	,    _A=False	,    _A="float32"	,    _A=False	,    _A=1_2_8	,    _A=6_4	,    _A=4	,    _A=4	,    _A=0.0_01	,    _A=0.0_01	,    _A="all"	,    _A=False	,    _A=False	,    _A=1.0	,    _A=0.2	,    _A=1	,    _A=0	,    _A=2	,    _A=False	,    **_A	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = vocab_size
											__lowerCAmelCase            = max_position_embeddings
											__lowerCAmelCase            = d_model
											__lowerCAmelCase            = encoder_ffn_dim
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = encoder_attention_heads
											__lowerCAmelCase            = decoder_ffn_dim
											__lowerCAmelCase            = decoder_layers
											__lowerCAmelCase            = decoder_attention_heads
											__lowerCAmelCase            = dropout
											__lowerCAmelCase            = attention_dropout
											__lowerCAmelCase            = activation_dropout
											__lowerCAmelCase            = activation_function
											__lowerCAmelCase            = init_std
											__lowerCAmelCase            = encoder_layerdrop
											__lowerCAmelCase            = decoder_layerdrop
											__lowerCAmelCase            = use_cache
											__lowerCAmelCase            = encoder_layers
											__lowerCAmelCase            = scale_embedding  # scale factor will be sqrt(d_model) if True
											__lowerCAmelCase            = router_z_loss_coef
											__lowerCAmelCase            = router_aux_loss_coef
											__lowerCAmelCase            = decoder_sparse_step
											__lowerCAmelCase            = encoder_sparse_step
											__lowerCAmelCase            = num_experts
											__lowerCAmelCase            = expert_capacity
											__lowerCAmelCase            = router_bias
											if router_dtype not in ["float32", "float16", "bfloat16"]:
																raise ValueError(f"""`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}"""    )
											__lowerCAmelCase            = router_dtype

											__lowerCAmelCase            = router_ignore_padding_tokens
											__lowerCAmelCase            = batch_prioritized_routing
											__lowerCAmelCase            = second_expert_policy
											__lowerCAmelCase            = normalize_router_prob_before_dropping
											__lowerCAmelCase            = moe_eval_capacity_token_fraction
											__lowerCAmelCase            = moe_token_dropout
											__lowerCAmelCase            = output_router_logits
											super().__init__(
											    pad_token_id=_A	,    bos_token_id=_A	,    eos_token_id=_A	,    is_encoder_decoder=_A	,    decoder_start_token_id=_A	,    **_A	,    )




