
"""simple docstring"""







import unittest

from transformers import DebertaVaConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           DebertaVaForMaskedLM,
           DebertaVaForMultipleChoice,
           DebertaVaForQuestionAnswering,
           DebertaVaForSequenceClassification,
           DebertaVaForTokenClassification,
           DebertaVaModel,
       )
       from transformers.models.deberta_va.modeling_deberta_va import DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




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



      """simple docstring"""

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




 def 	a   (			self							:			int			)->       Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	ids_tensor([self.batch_size, self.seq_length]  , self.vocab_size			)

      UpperCAmelCase_    :     Any            =	None
      if self.use_input_mask:
           UpperCAmelCase_    :     Tuple            =	ids_tensor([self.batch_size, self.seq_length]  , vocab_size=2			)

      UpperCAmelCase_    :     Optional[Any]            =	None
      if self.use_token_type_ids:
           UpperCAmelCase_    :     str            =	ids_tensor([self.batch_size, self.seq_length]  , self.type_vocab_size			)

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

      UpperCAmelCase_    :     Any            =	self.get_config()

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




 def 	a   (			self							:			List[Any]			)->       List[Any]:



      """simple docstring"""

      return 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  , initializer_range=self.initializer_range  , relative_attention=self.relative_attention  , position_biased_input=self.position_biased_input  , pos_att_type=self.pos_att_type  , )




 def 	a   (			self							:			Optional[int]  , a_							:			int			)->       Optional[Any]:



      """simple docstring"""

      self.parent.assertListEqual(list(result.loss.size()			)  , []			)




 def 	a   (			self							:			str  , a_							:			Tuple  , a_							:			List[str]  , a_							:			Dict  , a_							:			Union[str, Any]  , a_							:			Dict  , a_							:			List[str]  , a_							:			List[str]			)->       int:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	DebertaVaModel(config=a_			)
      model.to(a_			)
      model.eval()
      UpperCAmelCase_    :     Any            =	model(a_  , attention_mask=a_  , token_type_ids=a_			)[0]
      UpperCAmelCase_    :     Optional[Any]            =	model(a_  , token_type_ids=a_			)[0]
      UpperCAmelCase_    :     Dict            =	model(a_			)[0]

      self.parent.assertListEqual(list(sequence_output.size()			)  , [self.batch_size, self.seq_length, self.hidden_size]			)




 def 	a   (			self							:			int  , a_							:			int  , a_							:			int  , a_							:			Union[str, Any]  , a_							:			Tuple  , a_							:			List[Any]  , a_							:			Any  , a_							:			Tuple			)->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[Any]            =	DebertaVaForMaskedLM(config=a_			)
      model.to(a_			)
      model.eval()
      UpperCAmelCase_    :     int            =	model(a_  , attention_mask=a_  , token_type_ids=a_  , labels=a_			)

      self.parent.assertEqual(result.logits.shape  , (self.batch_size, self.seq_length, self.vocab_size)			)




 def 	a   (			self							:			Dict  , a_							:			int  , a_							:			Optional[int]  , a_							:			Union[str, Any]  , a_							:			int  , a_							:			int  , a_							:			int  , a_							:			Tuple			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[Any]            =	self.num_labels
      UpperCAmelCase_    :     Union[str, Any]            =	DebertaVaForSequenceClassification(a_			)
      model.to(a_			)
      model.eval()
      UpperCAmelCase_    :     Optional[Any]            =	model(a_  , attention_mask=a_  , token_type_ids=a_  , labels=a_			)
      self.parent.assertListEqual(list(result.logits.size()			)  , [self.batch_size, self.num_labels]			)
      self.check_loss_output(a_			)




 def 	a   (			self							:			str  , a_							:			str  , a_							:			Optional[int]  , a_							:			Optional[Any]  , a_							:			List[str]  , a_							:			List[str]  , a_							:			str  , a_							:			Any			)->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     str            =	self.num_labels
      UpperCAmelCase_    :     Dict            =	DebertaVaForTokenClassification(config=a_			)
      model.to(a_			)
      model.eval()
      UpperCAmelCase_    :     Any            =	model(a_  , attention_mask=a_  , token_type_ids=a_  , labels=a_			)
      self.parent.assertEqual(result.logits.shape  , (self.batch_size, self.seq_length, self.num_labels)			)




 def 	a   (			self							:			Any  , a_							:			Dict  , a_							:			int  , a_							:			str  , a_							:			Union[str, Any]  , a_							:			Optional[Any]  , a_							:			Any  , a_							:			Union[str, Any]			)->       Optional[int]:



      """simple docstring"""

      UpperCAmelCase_    :     List[str]            =	DebertaVaForQuestionAnswering(config=a_			)
      model.to(a_			)
      model.eval()
      UpperCAmelCase_    :     Any            =	model(
          a_  , attention_mask=a_  , token_type_ids=a_  , start_positions=a_  , end_positions=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							:			Dict  , a_							:			Tuple  , a_							:			Tuple  , a_							:			str  , a_							:			List[Any]  , a_							:			List[str]  , a_							:			Tuple  , a_							:			Union[str, Any]			)->       str:



      """simple docstring"""

      UpperCAmelCase_    :     List[str]            =	DebertaVaForMultipleChoice(config=a_			)
      model.to(a_			)
      model.eval()
      UpperCAmelCase_    :     Optional[Any]            =	input_ids.unsqueeze(1			).expand(-1  , self.num_choices  , -1			).contiguous()
      UpperCAmelCase_    :     Optional[int]            =	token_type_ids.unsqueeze(1			).expand(-1  , self.num_choices  , -1			).contiguous()
      UpperCAmelCase_    :     List[Any]            =	input_mask.unsqueeze(1			).expand(-1  , self.num_choices  , -1			).contiguous()
      UpperCAmelCase_    :     List[str]            =	model(
          a_  , attention_mask=a_  , token_type_ids=a_  , labels=a_  , )
      self.parent.assertEqual(result.logits.shape  , (self.batch_size, self.num_choices)			)






 def 	a   (			self							:			List[str]			)->       Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_    :     int            =	self.prepare_config_and_inputs()
      (
          (
          UpperCAmelCase_
      )       ,(
          UpperCAmelCase_
      )       ,(
          UpperCAmelCase_
      )       ,(
          UpperCAmelCase_
      )       ,(
          UpperCAmelCase_
      )       ,(
          UpperCAmelCase_
      )       ,(
          UpperCAmelCase_
      )       ,
      )    :     List[str]            =	config_and_inputs
      UpperCAmelCase_    :     List[Any]            =	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
      return config, inputs_dict



@require_torch
class 		UpperCAmelCase_						(lowerCamelCase_	,			lowerCamelCase_	,			unittest.TestCase    ):
 """simple docstring"""




 UpperCamelCase_    :  str					 =     (
     (
         DebertaVaModel,
         DebertaVaForMaskedLM,
         DebertaVaForSequenceClassification,
         DebertaVaForTokenClassification,
         DebertaVaForQuestionAnswering,
         DebertaVaForMultipleChoice,
     )
     if is_torch_available()
     else ()
 )
 UpperCamelCase_    :  List[str]					 =     (
     {
         """feature-extraction""": DebertaVaModel,
         """fill-mask""": DebertaVaForMaskedLM,
         """question-answering""": DebertaVaForQuestionAnswering,
         """text-classification""": DebertaVaForSequenceClassification,
         """token-classification""": DebertaVaForTokenClassification,
         """zero-shot""": DebertaVaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )

 UpperCamelCase_    :  int					 =     True
 UpperCamelCase_    :  int					 =     False
 UpperCamelCase_    :  str					 =     False
 UpperCamelCase_    :  List[str]					 =     False
 UpperCamelCase_    :  str					 =     False




 def 	a   (			self							:			Union[str, Any]			)->       int:



      """simple docstring"""

      UpperCAmelCase_    :     Union[str, Any]            =	DebertaVaModelTester(self			)
      UpperCAmelCase_    :     Any            =	ConfigTester(self  , config_class=a_  , hidden_size=37			)




 def 	a   (			self							:			str			)->       Any:



      """simple docstring"""

      self.config_tester.run_common_tests()




 def 	a   (			self							:			Optional[int]			)->       Dict:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_deberta_model(*a_			)




 def 	a   (			self							:			Any			)->       Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_deberta_for_sequence_classification(*a_			)




 def 	a   (			self							:			Optional[int]			)->       str:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_deberta_for_masked_lm(*a_			)




 def 	a   (			self							:			Any			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Any            =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_deberta_for_question_answering(*a_			)




 def 	a   (			self							:			int			)->       Tuple:



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_deberta_for_token_classification(*a_			)




 def 	a   (			self							:			Optional[Any]			)->       Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_deberta_for_multiple_choice(*a_			)






 @slow
 def 	a   (			self							:			List[str]			)->       Optional[Any]:



      """simple docstring"""

      for model_name in DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           UpperCAmelCase_    :     List[str]            =	DebertaVaModel.from_pretrained(a_			)
           self.assertIsNotNone(a_			)



@require_torch
@require_sentencepiece
@require_tokenizers
class 		UpperCAmelCase_						(unittest.TestCase    ):
 """simple docstring"""




 @unittest.skip(reason="""Model not available yet"""			)
 def 	a   (			self							:			Union[str, Any]			)->       Tuple:



      """simple docstring"""

      pass






 @slow
 def 	a   (			self							:			Tuple			)->       Dict:



      """simple docstring"""

      UpperCAmelCase_    :     int            =	DebertaVaModel.from_pretrained("""microsoft/deberta-v2-xlarge"""			)

      UpperCAmelCase_    :     Tuple            =	torch.tensor([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]			)
      UpperCAmelCase_    :     List[Any]            =	torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
      with torch.no_grad():
           UpperCAmelCase_    :     str            =	model(a_  , attention_mask=a_			)[0]
      # compare the actual values for a slice.
      UpperCAmelCase_    :     Tuple            =	torch.tensor(
          [[[0.2356, 0.1948, 0.0369], [-0.1063, 0.3586, -0.5152], [-0.6399, -0.0259, -0.2525]]]			)
      self.assertTrue(torch.allclose(output[:, 1:4, 1:4]  , a_  , atol=1E-4			)  , f'''{output[:, 1:4, 1:4]}'''			)



"""simple docstring"""







import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


lowercase_         =       logging.getLogger(__name__)



class 		UpperCAmelCase_						:
 """simple docstring"""




 def __init__(			self							:			Dict			)->       Optional[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	False




 def 	a   (			self							:			Optional[int]  , a_							:			Union[str, Any]  , a_							:			Tuple  , a_							:			Dict  , a_							:			Tuple			)->       Any:



      """simple docstring"""

      if not self.initialized:
           UpperCAmelCase_    :     Any            =	RagRetriever(
               a_  , question_encoder_tokenizer=a_  , generator_tokenizer=a_  , index=a_  , init_retrieval=a_  , )
           UpperCAmelCase_    :     Any            =	True




 def 	a   (			self							:			Dict			)->       Union[str, Any]:



      """simple docstring"""

      self.retriever.index.init_index()






 def 	a   (			self							:			Any  , a_							:			Any  , a_							:			Dict			)->       str:



      """simple docstring"""

      UpperCAmelCase_       ,UpperCAmelCase_    :     List[Any]            =	self.retriever._main_retrieve(a_  , a_			)
      return doc_ids, retrieved_doc_embeds



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def __init__(			self							:			int  , a_							:			Optional[int]  , a_							:			List[Any]  , a_							:			Optional[int]  , a_							:			List[Any]  , a_							:			Union[str, Any]=None			)->       Union[str, Any]:



      """simple docstring"""

      if index is not None and index.is_initialized() and len(a_			) > 0:
           raise ValueError(
               """When using Ray for distributed fine-tuning, """
               """you'll need to provide the paths instead, """
               """as the dataset and the index are loaded """
               """separately. More info in examples/rag/use_own_knowledge_dataset.py """			)
      super().__init__(
          a_  , question_encoder_tokenizer=a_  , generator_tokenizer=a_  , index=a_  , init_retrieval=a_  , )
      UpperCAmelCase_    :     Union[str, Any]            =	retrieval_workers
      if len(self.retrieval_workers			) > 0:
           ray.get(
               [
                   worker.create_rag_retriever.remote(a_  , a_  , a_  , a_			)
                   for worker in self.retrieval_workers
               ]			)




 def 	a   (			self							:			List[str]			)->       str:



      """simple docstring"""

      logger.info("""initializing retrieval"""			)

      if len(self.retrieval_workers			) > 0:
           ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]			)
      else:
           # Non-distributed training. Load index into this same process.
           self.index.init_index()




 def 	a   (			self							:			Optional[Any]  , a_							:			Optional[Any]  , a_							:			List[str]			)->       Tuple:



      """simple docstring"""

      if len(self.retrieval_workers			) > 0:
           # Select a random retrieval actor.
           UpperCAmelCase_    :     str            =	self.retrieval_workers[random.randint(0  , len(self.retrieval_workers			) - 1			)]
           UpperCAmelCase_       ,UpperCAmelCase_    :     List[str]            =	ray.get(random_worker.retrieve.remote(a_  , a_			)			)
      else:
           UpperCAmelCase_       ,UpperCAmelCase_    :     Tuple            =	self._main_retrieve(a_  , a_			)
      return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(a_			)




 @classmethod
 def 	a   (			cls							:			Dict  , a_							:			List[Any]  , a_							:			Union[str, Any]=None  , **a_							:			int			)->       str:



      """simple docstring"""

      return super(a_  , cls			).get_tokenizers(a_  , a_  , **a_			)






 @classmethod
 def 	a   (			cls							:			Tuple  , a_							:			Tuple  , a_							:			List[str]  , a_							:			Optional[int]=None  , **a_							:			int			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     List[str]            =	kwargs.pop("""config"""  , a_			) or RagConfig.from_pretrained(a_  , **a_			)
      UpperCAmelCase_    :     Optional[int]            =	RagTokenizer.from_pretrained(a_  , config=a_			)
      UpperCAmelCase_    :     Optional[Any]            =	rag_tokenizer.question_encoder
      UpperCAmelCase_    :     List[str]            =	rag_tokenizer.generator
      if indexed_dataset is not None:
           UpperCAmelCase_    :     Dict            =	"""custom"""
           UpperCAmelCase_    :     List[Any]            =	CustomHFIndex(config.retrieval_vector_size  , a_			)
      else:
           UpperCAmelCase_    :     int            =	cls._build_index(a_			)
      return cls(
          a_  , question_encoder_tokenizer=a_  , generator_tokenizer=a_  , retrieval_workers=a_  , index=a_  , )

