

'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
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 (
     TF_MODEL_FOR_PRETRAINING_MAPPING,
     TFMobileBertForMaskedLM,
     TFMobileBertForMultipleChoice,
     TFMobileBertForNextSentencePrediction,
     TFMobileBertForPreTraining,
     TFMobileBertForQuestionAnswering,
     TFMobileBertForSequenceClassification,
     TFMobileBertForTokenClassification,
     TFMobileBertModel,
 )




@require_tf
class        A_ (				_snake_case       ,   _snake_case       ,   unittest.TestCase							):






     '''simple docstring'''



     UpperCAmelCase_						:	Any  		=  (
         (
             TFMobileBertModel,
             TFMobileBertForMaskedLM,
             TFMobileBertForNextSentencePrediction,
             TFMobileBertForPreTraining,
             TFMobileBertForQuestionAnswering,
             TFMobileBertForSequenceClassification,
             TFMobileBertForTokenClassification,
             TFMobileBertForMultipleChoice,
         )
         if is_tf_available()
         else ()
     )
     UpperCAmelCase_						:	Optional[int]  		=  (
         {
             'feature-extraction': TFMobileBertModel,
             'fill-mask': TFMobileBertForMaskedLM,
             'question-answering': TFMobileBertForQuestionAnswering,
             'text-classification': TFMobileBertForSequenceClassification,
             'token-classification': TFMobileBertForTokenClassification,
             'zero-shot': TFMobileBertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     UpperCAmelCase_						:	str  		=  False
     UpperCAmelCase_						:	List[str]  		=  False


     def    UpperCAmelCase_       (		self :     Dict     ,			lowercase_ :     Union[str, Any]     ,			lowercase_ :     List[str]     ,			lowercase_ :     List[str]=False       )						->	Union[str, Any]:
      UpperCAmelCase       :	Union[str, Any]											= super()._prepare_for_class(__lowerCAmelCase     ,			__lowerCAmelCase     ,			return_labels=__lowerCAmelCase       )

      if return_labels:
       if model_class in get_values(__lowerCAmelCase       ):
        UpperCAmelCase       :	str											= tf.zeros(self.model_tester.batch_size     ,			dtype=tf.intaa       )

      return inputs_dict




     class        A_ (				_snake_case							):






          '''simple docstring'''



          def __init__(		self :     Any     ,			lowercase_ :     Tuple     ,			lowercase_ :     Optional[int]=13     ,			lowercase_ :     List[Any]=7     ,			lowercase_ :     Union[str, Any]=True     ,			lowercase_ :     Optional[int]=True     ,			lowercase_ :     Union[str, Any]=True     ,			lowercase_ :     Optional[Any]=True     ,			lowercase_ :     List[Any]=99     ,			lowercase_ :     List[Any]=32     ,			lowercase_ :     List[Any]=32     ,			lowercase_ :     Dict=2     ,			lowercase_ :     Dict=4     ,			lowercase_ :     Optional[int]=37     ,			lowercase_ :     List[Any]="gelu"     ,			lowercase_ :     List[Any]=0.1     ,			lowercase_ :     Optional[Any]=0.1     ,			lowercase_ :     str=512     ,			lowercase_ :     Union[str, Any]=16     ,			lowercase_ :     Dict=2     ,			lowercase_ :     Union[str, Any]=0.02     ,			lowercase_ :     List[Any]=3     ,			lowercase_ :     Tuple=4     ,			lowercase_ :     List[Any]=None     ,			)						->	Optional[int]:
           UpperCAmelCase       :	List[Any]											= parent
           UpperCAmelCase       :	Any											= batch_size
           UpperCAmelCase       :	Optional[int]											= seq_length
           UpperCAmelCase       :	List[str]											= is_training
           UpperCAmelCase       :	Any											= use_input_mask
           UpperCAmelCase       :	Tuple											= use_token_type_ids
           UpperCAmelCase       :	List[str]											= use_labels
           UpperCAmelCase       :	str											= vocab_size
           UpperCAmelCase       :	str											= hidden_size
           UpperCAmelCase       :	int											= num_hidden_layers
           UpperCAmelCase       :	Optional[Any]											= num_attention_heads
           UpperCAmelCase       :	int											= intermediate_size
           UpperCAmelCase       :	Optional[Any]											= hidden_act
           UpperCAmelCase       :	Any											= hidden_dropout_prob
           UpperCAmelCase       :	List[str]											= attention_probs_dropout_prob
           UpperCAmelCase       :	List[str]											= max_position_embeddings
           UpperCAmelCase       :	int											= type_vocab_size
           UpperCAmelCase       :	int											= type_sequence_label_size
           UpperCAmelCase       :	str											= initializer_range
           UpperCAmelCase       :	int											= num_labels
           UpperCAmelCase       :	List[Any]											= num_choices
           UpperCAmelCase       :	List[str]											= scope
           UpperCAmelCase       :	Optional[Any]											= embedding_size


          def    UpperCAmelCase_       (		self :     str       )						->	str:
           UpperCAmelCase       :	Optional[Any]											= ids_tensor([self.batch_size, self.seq_length]     ,			self.vocab_size       )

           UpperCAmelCase       :	Tuple											= None
           if self.use_input_mask:
            UpperCAmelCase       :	Dict											= random_attention_mask([self.batch_size, self.seq_length]       )

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

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

           UpperCAmelCase       :	List[str]											= MobileBertConfig(
               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     ,			embedding_size=self.embedding_size     ,			)

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


          def    UpperCAmelCase_       (		self :     Dict     ,			lowercase_ :     Union[str, Any]     ,			lowercase_ :     Any     ,			lowercase_ :     Optional[Any]     ,			lowercase_ :     Tuple     ,			lowercase_ :     Optional[int]     ,			lowercase_ :     Optional[Any]     ,			lowercase_ :     str       )						->	Union[str, Any]:
           UpperCAmelCase       :	str											= TFMobileBertModel(config=__lowerCAmelCase       )
           UpperCAmelCase       :	int											= {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
           UpperCAmelCase       :	Any											= model(__lowerCAmelCase       )

           UpperCAmelCase       :	str											= [input_ids, input_mask]
           UpperCAmelCase       :	Union[str, Any]											= model(__lowerCAmelCase       )

           UpperCAmelCase       :	Any											= model(__lowerCAmelCase       )

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


          def    UpperCAmelCase_       (		self :     Optional[Any]     ,			lowercase_ :     Union[str, Any]     ,			lowercase_ :     Optional[int]     ,			lowercase_ :     List[Any]     ,			lowercase_ :     List[str]     ,			lowercase_ :     Optional[int]     ,			lowercase_ :     Tuple     ,			lowercase_ :     Any       )						->	Union[str, Any]:
           UpperCAmelCase       :	List[Any]											= TFMobileBertForMaskedLM(config=__lowerCAmelCase       )
           UpperCAmelCase       :	Tuple											= {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
           UpperCAmelCase       :	Union[str, Any]											= model(__lowerCAmelCase       )
           self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.vocab_size)       )


          def    UpperCAmelCase_       (		self :     List[str]     ,			lowercase_ :     List[Any]     ,			lowercase_ :     Tuple     ,			lowercase_ :     Tuple     ,			lowercase_ :     Tuple     ,			lowercase_ :     int     ,			lowercase_ :     Tuple     ,			lowercase_ :     Optional[int]       )						->	Optional[int]:
           UpperCAmelCase       :	Optional[Any]											= TFMobileBertForNextSentencePrediction(config=__lowerCAmelCase       )
           UpperCAmelCase       :	Optional[Any]											= {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
           UpperCAmelCase       :	Dict											= model(__lowerCAmelCase       )
           self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, 2)       )


          def    UpperCAmelCase_       (		self :     Tuple     ,			lowercase_ :     Tuple     ,			lowercase_ :     Dict     ,			lowercase_ :     List[Any]     ,			lowercase_ :     Any     ,			lowercase_ :     List[str]     ,			lowercase_ :     List[Any]     ,			lowercase_ :     List[Any]       )						->	Optional[int]:
           UpperCAmelCase       :	str											= TFMobileBertForPreTraining(config=__lowerCAmelCase       )
           UpperCAmelCase       :	List[Any]											= {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
           UpperCAmelCase       :	Optional[Any]											= model(__lowerCAmelCase       )
           self.parent.assertEqual(
               result.prediction_logits.shape     ,			(self.batch_size, self.seq_length, self.vocab_size)       )
           self.parent.assertEqual(result.seq_relationship_logits.shape     ,			(self.batch_size, 2)       )


          def    UpperCAmelCase_       (		self :     Optional[int]     ,			lowercase_ :     List[Any]     ,			lowercase_ :     Optional[int]     ,			lowercase_ :     str     ,			lowercase_ :     Any     ,			lowercase_ :     Any     ,			lowercase_ :     int     ,			lowercase_ :     str       )						->	Optional[int]:
           UpperCAmelCase       :	int											= self.num_labels
           UpperCAmelCase       :	Optional[Any]											= TFMobileBertForSequenceClassification(config=__lowerCAmelCase       )
           UpperCAmelCase       :	List[Any]											= {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
           UpperCAmelCase       :	Any											= model(__lowerCAmelCase       )
           self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_labels)       )


          def    UpperCAmelCase_       (		self :     List[str]     ,			lowercase_ :     Tuple     ,			lowercase_ :     int     ,			lowercase_ :     Any     ,			lowercase_ :     int     ,			lowercase_ :     List[str]     ,			lowercase_ :     Optional[Any]     ,			lowercase_ :     Any       )						->	Union[str, Any]:
           UpperCAmelCase       :	Optional[int]											= self.num_choices
           UpperCAmelCase       :	Optional[Any]											= TFMobileBertForMultipleChoice(config=__lowerCAmelCase       )
           UpperCAmelCase       :	List[Any]											= tf.tile(tf.expand_dims(__lowerCAmelCase     ,			1       )     ,			(1, self.num_choices, 1)       )
           UpperCAmelCase       :	str											= tf.tile(tf.expand_dims(__lowerCAmelCase     ,			1       )     ,			(1, self.num_choices, 1)       )
           UpperCAmelCase       :	List[Any]											= tf.tile(tf.expand_dims(__lowerCAmelCase     ,			1       )     ,			(1, self.num_choices, 1)       )
           UpperCAmelCase       :	str											= {
               'input_ids': multiple_choice_inputs_ids,
               'attention_mask': multiple_choice_input_mask,
               'token_type_ids': multiple_choice_token_type_ids,
           }
           UpperCAmelCase       :	Dict											= model(__lowerCAmelCase       )
           self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_choices)       )


          def    UpperCAmelCase_       (		self :     int     ,			lowercase_ :     str     ,			lowercase_ :     int     ,			lowercase_ :     Dict     ,			lowercase_ :     Any     ,			lowercase_ :     int     ,			lowercase_ :     Any     ,			lowercase_ :     int       )						->	Optional[Any]:
           UpperCAmelCase       :	List[Any]											= self.num_labels
           UpperCAmelCase       :	Any											= TFMobileBertForTokenClassification(config=__lowerCAmelCase       )
           UpperCAmelCase       :	Dict											= {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
           UpperCAmelCase       :	int											= model(__lowerCAmelCase       )
           self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.num_labels)       )


          def    UpperCAmelCase_       (		self :     Optional[Any]     ,			lowercase_ :     List[str]     ,			lowercase_ :     int     ,			lowercase_ :     Any     ,			lowercase_ :     str     ,			lowercase_ :     List[Any]     ,			lowercase_ :     Optional[int]     ,			lowercase_ :     Tuple       )						->	int:
           UpperCAmelCase       :	Dict											= TFMobileBertForQuestionAnswering(config=__lowerCAmelCase       )
           UpperCAmelCase       :	int											= {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
           UpperCAmelCase       :	Optional[Any]											= 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    UpperCAmelCase_       (		self :     Optional[Any]       )						->	Any:
           UpperCAmelCase       :	str											= self.prepare_config_and_inputs()
           (
               (
               UpperCAmelCase							
           )					, (
               UpperCAmelCase							
           )					, (
               UpperCAmelCase							
           )					, (
               UpperCAmelCase							
           )					, (
               UpperCAmelCase							
           )					, (
               UpperCAmelCase							
           )					, (
               UpperCAmelCase							
           )					, 							
           )       :	List[str]											= config_and_inputs
           UpperCAmelCase       :	Optional[Any]											= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
           return config, inputs_dict



     def    UpperCAmelCase_       (		self :     Optional[Any]       )						->	Tuple:
      UpperCAmelCase       :	Dict											= TFMobileBertModelTest.TFMobileBertModelTester(self       )
      UpperCAmelCase       :	Dict											= ConfigTester(self     ,			config_class=__lowerCAmelCase     ,			hidden_size=37       )


     def    UpperCAmelCase_       (		self :     Union[str, Any]       )						->	int:
      self.config_tester.run_common_tests()


     def    UpperCAmelCase_       (		self :     List[Any]       )						->	Optional[Any]:
      UpperCAmelCase       :	int											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_model(*__lowerCAmelCase       )


     def    UpperCAmelCase_       (		self :     List[Any]       )						->	Optional[Any]:
      UpperCAmelCase       :	Any											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_masked_lm(*__lowerCAmelCase       )


     def    UpperCAmelCase_       (		self :     Dict       )						->	Tuple:
      UpperCAmelCase       :	Optional[int]											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_multiple_choice(*__lowerCAmelCase       )


     def    UpperCAmelCase_       (		self :     Tuple       )						->	str:
      UpperCAmelCase       :	Union[str, Any]											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*__lowerCAmelCase       )


     def    UpperCAmelCase_       (		self :     Optional[int]       )						->	List[str]:
      UpperCAmelCase       :	Optional[int]											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_pretraining(*__lowerCAmelCase       )


     def    UpperCAmelCase_       (		self :     int       )						->	int:
      UpperCAmelCase       :	int											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_question_answering(*__lowerCAmelCase       )


     def    UpperCAmelCase_       (		self :     Union[str, Any]       )						->	Any:
      UpperCAmelCase       :	Any											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_sequence_classification(*__lowerCAmelCase       )


     def    UpperCAmelCase_       (		self :     Any       )						->	Union[str, Any]:
      UpperCAmelCase       :	int											= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_token_classification(*__lowerCAmelCase       )




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



      for model_name in ["google/mobilebert-uncased"]:
       UpperCAmelCase       :	int											= TFMobileBertModel.from_pretrained(__lowerCAmelCase       )
       self.assertIsNotNone(__lowerCAmelCase       )




@require_tf
class        A_ (				unittest.TestCase							):






     '''simple docstring'''



     @slow
     def    UpperCAmelCase_       (		self :     Any       )						->	Tuple:
      UpperCAmelCase       :	Optional[Any]											= TFMobileBertForPreTraining.from_pretrained('google/mobilebert-uncased'       )
      UpperCAmelCase       :	Optional[int]											= tf.constant([[0, 1, 2, 3, 4, 5]]       )
      UpperCAmelCase       :	Dict											= model(__lowerCAmelCase       )[0]

      UpperCAmelCase       :	List[str]											= [1, 6, 30_522]
      self.assertEqual(output.shape     ,			__lowerCAmelCase       )

      UpperCAmelCase       :	Optional[Any]											= tf.constant(
          [
              [
                  [-4.591_9547, -9.24_8295, -9.64_5256],
                  [-6.730_6175, -6.44_0284, -6.605_2837],
                  [-7.274_3506, -6.784_7915, -6.02_4673],
              ]
          ]       )
      tf.debugging.assert_near(output[:, :3, :3]     ,			__lowerCAmelCase     ,			atol=1E-4       )





"""simple docstring"""






import logging
import random

import ray

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


__lowerCAmelCase    :       int					     =logging.getLogger(__name__)





class  _A						:

      def __init__(						self						):







            """simple docstring"""



            lowercase														=	False

      def 				A__       (						self		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase						):







            """simple docstring"""



            if not self.initialized:
                  lowercase														=	RagRetriever(
                      __lowerCAmelCase		,       question_encoder_tokenizer=__lowerCAmelCase		,       generator_tokenizer=__lowerCAmelCase		,       index=__lowerCAmelCase		,       init_retrieval=__lowerCAmelCase		,       )
                  lowercase														=	True

      def 				A__       (						self						):







            """simple docstring"""



            self.retriever.index.init_index()







      def 				A__       (						self		,       __lowerCAmelCase		,       __lowerCAmelCase						):







            """simple docstring"""



            lowercase							,							lowercase														=	self.retriever._main_retrieve(__lowerCAmelCase		,       __lowerCAmelCase						)
            return doc_ids, retrieved_doc_embeds






class  _A						(       lowerCAmelCase	):

      def __init__(						self		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase=None						):







            """simple docstring"""



            if index is not None and index.is_initialized() and len(__lowerCAmelCase						) > 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__(
                __lowerCAmelCase		,       question_encoder_tokenizer=__lowerCAmelCase		,       generator_tokenizer=__lowerCAmelCase		,       index=__lowerCAmelCase		,       init_retrieval=__lowerCAmelCase		,       )
            lowercase														=	retrieval_workers
            if len(self.retrieval_workers						) > 0:
                  ray.get(
                      [
                          worker.create_rag_retriever.remote(__lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase						)
                          for worker in self.retrieval_workers
                      ]						)

      def 				A__       (						self						):







            """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		,       __lowerCAmelCase		,       __lowerCAmelCase						):







            """simple docstring"""



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

      @classmethod
      def 				A__       (						cls		,       __lowerCAmelCase		,       __lowerCAmelCase=None		,       **__lowerCAmelCase						):







            """simple docstring"""



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







      @classmethod
      def 				A__       (						cls		,       __lowerCAmelCase		,       __lowerCAmelCase		,       __lowerCAmelCase=None		,       **__lowerCAmelCase						):







            """simple docstring"""



            lowercase														=	kwargs.pop("""config"""		,       __lowerCAmelCase						) or RagConfig.from_pretrained(__lowerCAmelCase		,       **__lowerCAmelCase						)
            lowercase														=	RagTokenizer.from_pretrained(__lowerCAmelCase		,       config=__lowerCAmelCase						)
            lowercase														=	rag_tokenizer.question_encoder
            lowercase														=	rag_tokenizer.generator
            if indexed_dataset is not None:
                  lowercase														=	"""custom"""
                  lowercase														=	CustomHFIndex(config.retrieval_vector_size		,       __lowerCAmelCase						)
            else:
                  lowercase														=	cls._build_index(__lowerCAmelCase						)
            return cls(
                __lowerCAmelCase		,       question_encoder_tokenizer=__lowerCAmelCase		,       generator_tokenizer=__lowerCAmelCase		,       retrieval_workers=__lowerCAmelCase		,       index=__lowerCAmelCase		,       )



