






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

				"""simple docstring"""




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

								'''simple docstring'''
								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__    =			relative_attention
								lowercase__    =			position_biased_input
								lowercase__    =			pos_att_type
								lowercase__    =			scope




				def 				lowercase__   (     self						:    List[str]  ):

								'''simple docstring'''
								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
								if self.use_token_type_ids:
												lowercase__    =			ids_tensor([self.batch_size, self.seq_length],  self.type_vocab_size  )

								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__    =			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=lowerCamelCase,  )

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




				def 				lowercase__   (     self						:    List[Any],  lowerCamelCase						:    Optional[Any],  lowerCamelCase						:    Optional[int],  lowerCamelCase						:    Dict,  lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Tuple,  lowerCamelCase						:    int,  lowerCamelCase						:    Union[str, Any]  ):

								'''simple docstring'''
								lowercase__    =			TFDebertaVaModel(config=lowerCamelCase  )
								lowercase__    =			{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

								lowercase__    =			[input_ids, input_mask]
								lowercase__    =			model(lowerCamelCase  )

								lowercase__    =			model(lowerCamelCase  )

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




				def 				lowercase__   (     self						:    Any,  lowerCamelCase						:    Optional[int],  lowerCamelCase						:    Optional[Any],  lowerCamelCase						:    Tuple,  lowerCamelCase						:    Dict,  lowerCamelCase						:    str,  lowerCamelCase						:    int,  lowerCamelCase						:    Union[str, Any]  ):

								'''simple docstring'''
								lowercase__    =			TFDebertaVaForMaskedLM(config=lowerCamelCase  )
								lowercase__    =			{
								    '''input_ids''': input_ids,
								    '''attention_mask''': input_mask,
								    '''token_type_ids''': token_type_ids,
								}
								lowercase__    =			model(lowerCamelCase  )
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.seq_length, self.vocab_size)  )




				def 				lowercase__   (     self						:    Dict,  lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Dict,  lowerCamelCase						:    int,  lowerCamelCase						:    int,  lowerCamelCase						:    Any,  lowerCamelCase						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			self.num_labels
								lowercase__    =			TFDebertaVaForSequenceClassification(config=lowerCamelCase  )
								lowercase__    =			{
								    '''input_ids''': input_ids,
								    '''attention_mask''': input_mask,
								    '''token_type_ids''': token_type_ids,
								}

								lowercase__    =			model(lowerCamelCase  )
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.num_labels)  )




				def 				lowercase__   (     self						:    Union[str, Any],  lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Any,  lowerCamelCase						:    Any,  lowerCamelCase						:    int,  lowerCamelCase						:    Optional[Any],  lowerCamelCase						:    Tuple,  lowerCamelCase						:    Any  ):

								'''simple docstring'''
								lowercase__    =			self.num_labels
								lowercase__    =			TFDebertaVaForTokenClassification(config=lowerCamelCase  )
								lowercase__    =			{
								    '''input_ids''': input_ids,
								    '''attention_mask''': input_mask,
								    '''token_type_ids''': token_type_ids,
								}
								lowercase__    =			model(lowerCamelCase  )
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.seq_length, self.num_labels)  )




				def 				lowercase__   (     self						:    List[Any],  lowerCamelCase						:    int,  lowerCamelCase						:    Any,  lowerCamelCase						:    Optional[int],  lowerCamelCase						:    Optional[int],  lowerCamelCase						:    List[str],  lowerCamelCase						:    str,  lowerCamelCase						:    Union[str, Any]  ):

								'''simple docstring'''
								lowercase__    =			TFDebertaVaForQuestionAnswering(config=lowerCamelCase  )
								lowercase__    =			{
								    '''input_ids''': input_ids,
								    '''attention_mask''': input_mask,
								    '''token_type_ids''': token_type_ids,
								}

								lowercase__    =			model(lowerCamelCase  )
								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 				lowercase__   (     self						:    Dict  ):

								'''simple docstring'''
								lowercase__    =			self.prepare_config_and_inputs()
								(
								    (
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						
								)    =			config_and_inputs
								lowercase__    =			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
								return config, inputs_dict



@require_tf
class       _UpperCAmelCase					(       A__		,A__		,unittest.TestCase	):

				"""simple docstring"""




				lowercase__											=       (
				    (
				        TFDebertaVaModel,
				        TFDebertaVaForMaskedLM,
				        TFDebertaVaForQuestionAnswering,
				        TFDebertaVaForSequenceClassification,
				        TFDebertaVaForTokenClassification,
				    )
				    if is_tf_available()
				    else ()
				)
				lowercase__											=       (
				    {
				        """feature-extraction""": TFDebertaVaModel,
				        """fill-mask""": TFDebertaVaForMaskedLM,
				        """question-answering""": TFDebertaVaForQuestionAnswering,
				        """text-classification""": TFDebertaVaForSequenceClassification,
				        """token-classification""": TFDebertaVaForTokenClassification,
				        """zero-shot""": TFDebertaVaForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)

				lowercase__											=       False
				lowercase__											=       False




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			TFDebertaVaModelTester(self  )
								lowercase__    =			ConfigTester(self,  config_class=lowerCamelCase,  hidden_size=37  )




				def 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								self.config_tester.run_common_tests()




				def 				lowercase__   (     self						:    str  ):

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase  )




				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase  )




				def 				lowercase__   (     self						:    Dict  ):

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*lowerCamelCase  )




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

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase  )




				def 				lowercase__   (     self						:    int  ):

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*lowerCamelCase  )




				@slow
				def 				lowercase__   (     self						:    int  ):

								'''simple docstring'''
								lowercase__    =			TFDebertaVaModel.from_pretrained('''kamalkraj/deberta-v2-xlarge'''  )
								self.assertIsNotNone(lowerCamelCase  )



@require_tf
class       _UpperCAmelCase					(       unittest.TestCase	):

				"""simple docstring"""




				@unittest.skip(reason='''Model not available yet'''  )
				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								pass




				@slow
				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								lowercase__    =			TFDebertaVaModel.from_pretrained('''kamalkraj/deberta-v2-xlarge'''  )
								lowercase__    =			tf.constant([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]  )
								lowercase__    =			tf.constant([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
								lowercase__    =			model(lowerCamelCase,  attention_mask=lowerCamelCase  )[0]

								lowercase__    =			tf.constant(
								    [[[0.2356, 0.1948, 0.0369], [-0.1063, 0.3586, -0.5152], [-0.6399, -0.0259, -0.2525]]]  )
								tf.debugging.assert_near(output[:, 1:4, 1:4],  lowerCamelCase,  atol=1E-4  )












from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map






def   a			(						lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			[]
				lowercase__    =			[]
				lowercase__    =			[]

				for rt in rc.restypes:
								lowercase__    =			rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
								restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names]				)
								lowercase__    =			{name: i for i, name in enumerate(lowerCamelCase_				)}
								restype_atomaa_to_atomaa_list.append(
								    [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types]				)

								restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names]				)

				# Add dummy mapping for restype 'UNK'
				restype_atomaa_to_atomaa_list.append([0] * 14				)
				restype_atomaa_to_atomaa_list.append([0] * 37				)
				restype_atomaa_mask_list.append([0.0] * 14				)

				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.intaa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.intaa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.floataa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			protein['''aatype'''].to(torch.long				)

				# create the mapping for (residx, atom14) --> atom37, i.e. an array
				# with shape (num_res, 14) containing the atom37 indices for this protein
				lowercase__    =			restype_atomaa_to_atomaa[protein_aatype]
				lowercase__    =			restype_atomaa_mask[protein_aatype]

				lowercase__    =			residx_atomaa_mask
				lowercase__    =			residx_atomaa_to_atomaa.long()

				# create the gather indices for mapping back
				lowercase__    =			restype_atomaa_to_atomaa[protein_aatype]
				lowercase__    =			residx_atomaa_to_atomaa.long()

				# create the corresponding mask
				lowercase__    =			torch.zeros([21, 37]		,   dtype=torch.floataa		,   device=protein['''aatype'''].device				)
				for restype, restype_letter in enumerate(rc.restypes				):
								lowercase__    =			rc.restype_atoa[restype_letter]
								lowercase__    =			rc.residue_atoms[restype_name]
								for atom_name in atom_names:
												lowercase__    =			rc.atom_order[atom_name]
												lowercase__    =			1

				lowercase__    =			restype_atomaa_mask[protein_aatype]
				lowercase__    =			residx_atomaa_mask

				return protein






def   a			(						lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			tree_map(lambda lowerCamelCase_				: torch.tensor(lowerCamelCase_		,   device=batch['''aatype'''].device				)		,   lowerCamelCase_		,   np.ndarray				)
				lowercase__    =			tensor_tree_map(lambda lowerCamelCase_				: np.array(lowerCamelCase_				)		,   make_atomaa_masks(lowerCamelCase_				)				)
				return out




