


'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import DebertaVaConfig, 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 (
       TFDebertaVaForMaskedLM,
       TFDebertaVaForQuestionAnswering,
       TFDebertaVaForSequenceClassification,
       TFDebertaVaForTokenClassification,
       TFDebertaVaModel,
   )





class 						__SCREAMING_SNAKE_CASE					:






     def __init__(			self     ,							__UpperCamelCase     ,							__UpperCamelCase=13     ,							__UpperCamelCase=7     ,							__UpperCamelCase=True     ,							__UpperCamelCase=True     ,							__UpperCamelCase=True     ,							__UpperCamelCase=True     ,							__UpperCamelCase=99     ,							__UpperCamelCase=32     ,							__UpperCamelCase=2     ,							__UpperCamelCase=4     ,							__UpperCamelCase=37     ,							__UpperCamelCase="gelu"     ,							__UpperCamelCase=0.1     ,							__UpperCamelCase=0.1     ,							__UpperCamelCase=512     ,							__UpperCamelCase=16     ,							__UpperCamelCase=2     ,							__UpperCamelCase=0.02     ,							__UpperCamelCase=False     ,							__UpperCamelCase=True     ,							__UpperCamelCase="None"     ,							__UpperCamelCase=3     ,							__UpperCamelCase=4     ,							__UpperCamelCase=None     ,							)       ->    Optional[Any]:
          _a						=		parent
          _a						=		batch_size
          _a						=		seq_length
          _a						=		is_training
          _a						=		use_input_mask
          _a						=		use_token_type_ids
          _a						=		use_labels
          _a						=		vocab_size
          _a						=		hidden_size
          _a						=		num_hidden_layers
          _a						=		num_attention_heads
          _a						=		intermediate_size
          _a						=		hidden_act
          _a						=		hidden_dropout_prob
          _a						=		attention_probs_dropout_prob
          _a						=		max_position_embeddings
          _a						=		type_vocab_size
          _a						=		type_sequence_label_size
          _a						=		initializer_range
          _a						=		num_labels
          _a						=		num_choices
          _a						=		relative_attention
          _a						=		position_biased_input
          _a						=		pos_att_type
          _a						=		scope






     def  a_					(			self		)       ->    int:
          _a						=		ids_tensor([self.batch_size, self.seq_length]     ,							self.vocab_size		)

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

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

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

          _a						=		DebertaVaConfig(
              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     ,							relative_attention=self.relative_attention     ,							position_biased_input=self.position_biased_input     ,							initializer_range=self.initializer_range     ,							return_dict=__UpperCamelCase     ,							)

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






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Tuple:
          _a						=		TFDebertaVaModel(config=__UpperCamelCase		)
          _a						=		{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

          _a						=		[input_ids, input_mask]
          _a						=		model(__UpperCamelCase		)

          _a						=		model(__UpperCamelCase		)

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






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Union[str, Any]:
          _a						=		TFDebertaVaForMaskedLM(config=__UpperCamelCase		)
          _a						=		{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          _a						=		model(__UpperCamelCase		)
          self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.seq_length, self.vocab_size)		)






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    int:
          _a						=		self.num_labels
          _a						=		TFDebertaVaForSequenceClassification(config=__UpperCamelCase		)
          _a						=		{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }

          _a						=		model(__UpperCamelCase		)
          self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.num_labels)		)






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    str:
          _a						=		self.num_labels
          _a						=		TFDebertaVaForTokenClassification(config=__UpperCamelCase		)
          _a						=		{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          _a						=		model(__UpperCamelCase		)
          self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.seq_length, self.num_labels)		)






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    int:
          _a						=		TFDebertaVaForQuestionAnswering(config=__UpperCamelCase		)
          _a						=		{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }

          _a						=		model(__UpperCamelCase		)
          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		)       ->    int:
          _a						=		self.prepare_config_and_inputs()
          (
              (
              _a			
          )			,			(
              _a			
          )			,			(
              _a			
          )			,			(
              _a			
          )			,			(
              _a			
          )			,			(
              _a			
          )			,			(
              _a			
          )			,						
          )						=		config_and_inputs
          _a						=		{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
          return config, inputs_dict







@require_tf
class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ,			lowerCamelCase__ ,			unittest.TestCase ):
     UpperCAmelCase =				(
         (
             TFDebertaVaModel,
             TFDebertaVaForMaskedLM,
             TFDebertaVaForQuestionAnswering,
             TFDebertaVaForSequenceClassification,
             TFDebertaVaForTokenClassification,
         )
         if is_tf_available()
         else ()
     )
     UpperCAmelCase =				(
         {
             '''feature-extraction''': TFDebertaVaModel,
             '''fill-mask''': TFDebertaVaForMaskedLM,
             '''question-answering''': TFDebertaVaForQuestionAnswering,
             '''text-classification''': TFDebertaVaForSequenceClassification,
             '''token-classification''': TFDebertaVaForTokenClassification,
             '''zero-shot''': TFDebertaVaForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )

     UpperCAmelCase =				False
     UpperCAmelCase =				False






     def  a_					(			self		)       ->    Union[str, Any]:
          _a						=		TFDebertaVaModelTester(self		)
          _a						=		ConfigTester(self     ,							config_class=__UpperCamelCase     ,							hidden_size=37		)






     def  a_					(			self		)       ->    Optional[int]:
          self.config_tester.run_common_tests()






     def  a_					(			self		)       ->    Tuple:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__UpperCamelCase		)






     def  a_					(			self		)       ->    int:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase		)






     def  a_					(			self		)       ->    List[str]:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase		)






     def  a_					(			self		)       ->    str:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase		)






     def  a_					(			self		)       ->    int:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase		)






     @slow
     def  a_					(			self		)       ->    Any:
          _a						=		TFDebertaVaModel.from_pretrained("kamalkraj/deberta-v2-xlarge"		)
          self.assertIsNotNone(__UpperCamelCase		)







@require_tf
class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):






     @unittest.skip(reason="Model not available yet"		)
     def  a_					(			self		)       ->    List[Any]:
          pass






     @slow
     def  a_					(			self		)       ->    str:
          _a						=		TFDebertaVaModel.from_pretrained("kamalkraj/deberta-v2-xlarge"		)
          _a						=		tf.constant([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]		)
          _a						=		tf.constant([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
          _a						=		model(__UpperCamelCase     ,							attention_mask=__UpperCamelCase		)[0]

          _a						=		tf.constant(
              [[[0.23_56, 0.19_48, 0.03_69], [-0.10_63, 0.35_86, -0.51_52], [-0.63_99, -0.02_59, -0.25_25]]]		)
          tf.debugging.assert_near(output[:, 1:4, 1:4]     ,							__UpperCamelCase     ,							atol=1e-4		)





'''simple docstring'''

import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
lowercase__				     =		object()

# For specifying empty leaf dict `{}`
lowercase__				     =		object()






def __UpperCamelCase   (			__lowerCamelCase		: Tuple      ,  __lowerCamelCase		: Tuple				)      ->     Union[str, Any]:







     '''simple docstring'''






     _a						=		tuple((re.compile(x + "$"				) for x in qs)				)
     for i in range(len(__lowerCamelCase				) - len(__lowerCamelCase				) + 1				):
          _a						=		[x.match(__lowerCamelCase				) for x, y in zip(__lowerCamelCase      ,  ks[i:]				)]
          if matches and all(__lowerCamelCase				):
               return True
     return False






def __UpperCamelCase   (			__lowerCamelCase		: int				)      ->     Union[str, Any]:







     '''simple docstring'''






     def replace(__lowerCamelCase		: Tuple      ,  __lowerCamelCase		: Any				):
          for rule, replacement in rules:
               if _match(__lowerCamelCase      ,  __lowerCamelCase				):
                    return replacement
          return val

     return replace






def __UpperCamelCase   (			)      ->     Tuple:







     '''simple docstring'''






     return [
         # embeddings
         (("transformer", "wpe", "embedding"), P("mp"      ,  __lowerCamelCase				)),
         (("transformer", "wte", "embedding"), P("mp"      ,  __lowerCamelCase				)),
         # atention
         (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__lowerCamelCase      ,  "mp"				)),
         (("attention", "out_proj", "kernel"), P("mp"      ,  __lowerCamelCase				)),
         (("attention", "out_proj", "bias"), None),
         # mlp
         (("mlp", "c_fc", "kernel"), P(__lowerCamelCase      ,  "mp"				)),
         (("mlp", "c_fc", "bias"), P("mp"				)),
         (("mlp", "c_proj", "kernel"), P("mp"      ,  __lowerCamelCase				)),
         (("mlp", "c_proj", "bias"), None),
         # layer norms
         ((r"ln_\d+", "bias"), None),
         ((r"\d+", r"ln_\d+", "scale"), None),
         (("ln_f", "bias"), None),
         (("ln_f", "scale"), None),
     ]






def __UpperCamelCase   (			__lowerCamelCase		: Dict				)      ->     Tuple:







     '''simple docstring'''






     _a						=		_get_partition_rules()
     _a						=		_replacement_rules(__lowerCamelCase				)
     _a						=		{k: _unmatched for k in flatten_dict(__lowerCamelCase				)}
     _a						=		{k: replace(__lowerCamelCase      ,  __lowerCamelCase				) for k, v in initd.items()}
     assert _unmatched not in result.values(), "Incomplete partition spec."
     return freeze(unflatten_dict(__lowerCamelCase				)				)

