





'''simple docstring'''




import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

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


if is_torch_available():
    import torch

    from transformers import (
        FlaubertForMultipleChoice,
        FlaubertForQuestionAnswering,
        FlaubertForQuestionAnsweringSimple,
        FlaubertForSequenceClassification,
        FlaubertForTokenClassification,
        FlaubertModel,
        FlaubertWithLMHeadModel,
    )
    from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST






class UpperCAmelCase							(		UpperCAmelCase_							):



    def __init__(  self     ,			__A     ,			__A=13     ,			__A=7     ,			__A=True     ,			__A=True     ,			__A=True     ,			__A=True     ,			__A=True     ,			__A=False     ,			__A=False     ,			__A=False     ,			__A=2     ,			__A=99     ,			__A=0     ,			__A=32     ,			__A=5     ,			__A=4     ,			__A=0.1     ,			__A=0.1     ,			__A=512     ,			__A=12     ,			__A=2     ,			__A=0.0_2     ,			__A=3     ,			__A=4     ,			__A="last"     ,			__A=None     ,			__A=None     ,			):
     __UpperCAmelCase					=   parent
     __UpperCAmelCase					=   batch_size
     __UpperCAmelCase					=   seq_length
     __UpperCAmelCase					=   is_training
     __UpperCAmelCase					=   use_input_lengths
     __UpperCAmelCase					=   use_token_type_ids
     __UpperCAmelCase					=   use_labels
     __UpperCAmelCase					=   gelu_activation
     __UpperCAmelCase					=   sinusoidal_embeddings
     __UpperCAmelCase					=   causal
     __UpperCAmelCase					=   asm
     __UpperCAmelCase					=   n_langs
     __UpperCAmelCase					=   vocab_size
     __UpperCAmelCase					=   n_special
     __UpperCAmelCase					=   hidden_size
     __UpperCAmelCase					=   num_hidden_layers
     __UpperCAmelCase					=   num_attention_heads
     __UpperCAmelCase					=   hidden_dropout_prob
     __UpperCAmelCase					=   attention_probs_dropout_prob
     __UpperCAmelCase					=   max_position_embeddings
     __UpperCAmelCase					=   type_vocab_size
     __UpperCAmelCase					=   type_sequence_label_size
     __UpperCAmelCase					=   initializer_range
     __UpperCAmelCase					=   num_labels
     __UpperCAmelCase					=   num_choices
     __UpperCAmelCase					=   summary_type
     __UpperCAmelCase					=   use_proj
     __UpperCAmelCase					=   scope



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   ids_tensor([self.batch_size, self.seq_length]     ,			self.vocab_size					)
     __UpperCAmelCase					=   random_attention_mask([self.batch_size, self.seq_length]					)

     __UpperCAmelCase					=   None
     if self.use_input_lengths:
      __UpperCAmelCase					=   (
          ids_tensor([self.batch_size]     ,			vocab_size=2					) + self.seq_length - 2
      )  # small variation of seq_length

     __UpperCAmelCase					=   None
     if self.use_token_type_ids:
      __UpperCAmelCase					=   ids_tensor([self.batch_size, self.seq_length]     ,			self.n_langs					)

     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None
     if self.use_labels:
      __UpperCAmelCase					=   ids_tensor([self.batch_size]     ,			self.type_sequence_label_size					)
      __UpperCAmelCase					=   ids_tensor([self.batch_size, self.seq_length]     ,			self.num_labels					)
      __UpperCAmelCase					=   ids_tensor([self.batch_size]     ,			2					).float()
      __UpperCAmelCase					=   ids_tensor([self.batch_size]     ,			self.num_choices					)

     __UpperCAmelCase					=   self.get_config()

     return (
         config,
         input_ids,
         token_type_ids,
         input_lengths,
         sequence_labels,
         token_labels,
         is_impossible_labels,
         choice_labels,
         input_mask,
     )



    def 					__lowerCamelCase						(  self					):
     return FlaubertConfig(
         vocab_size=self.vocab_size     ,			n_special=self.n_special     ,			emb_dim=self.hidden_size     ,			n_layers=self.num_hidden_layers     ,			n_heads=self.num_attention_heads     ,			dropout=self.hidden_dropout_prob     ,			attention_dropout=self.attention_probs_dropout_prob     ,			gelu_activation=self.gelu_activation     ,			sinusoidal_embeddings=self.sinusoidal_embeddings     ,			asm=self.asm     ,			causal=self.causal     ,			n_langs=self.n_langs     ,			max_position_embeddings=self.max_position_embeddings     ,			initializer_range=self.initializer_range     ,			summary_type=self.summary_type     ,			use_proj=self.use_proj     ,			)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			):
     __UpperCAmelCase					=   FlaubertModel(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(__A     ,			lengths=__A     ,			langs=__A					)
     __UpperCAmelCase					=   model(__A     ,			langs=__A					)
     __UpperCAmelCase					=   model(__A					)
     self.parent.assertEqual(result.last_hidden_state.shape     ,			(self.batch_size, self.seq_length, self.hidden_size)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			):
     __UpperCAmelCase					=   FlaubertWithLMHeadModel(__A					)
     model.to(__A					)
     model.eval()

     __UpperCAmelCase					=   model(__A     ,			token_type_ids=__A     ,			labels=__A					)
     self.parent.assertEqual(result.loss.shape     ,			()					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.vocab_size)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			):
     __UpperCAmelCase					=   FlaubertForQuestionAnsweringSimple(__A					)
     model.to(__A					)
     model.eval()

     __UpperCAmelCase					=   model(__A					)

     __UpperCAmelCase					=   model(__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 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			):
     __UpperCAmelCase					=   FlaubertForQuestionAnswering(__A					)
     model.to(__A					)
     model.eval()

     __UpperCAmelCase					=   model(__A					)

     __UpperCAmelCase					=   model(
         __A     ,			start_positions=__A     ,			end_positions=__A     ,			cls_index=__A     ,			is_impossible=__A     ,			p_mask=__A     ,			)

     __UpperCAmelCase					=   model(
         __A     ,			start_positions=__A     ,			end_positions=__A     ,			cls_index=__A     ,			is_impossible=__A     ,			)

     ((__UpperCAmelCase)					,    )					=   result_with_labels.to_tuple()

     __UpperCAmelCase					=   model(__A     ,			start_positions=__A     ,			end_positions=__A					)

     ((__UpperCAmelCase)					,    )					=   result_with_labels.to_tuple()

     self.parent.assertEqual(result_with_labels.loss.shape     ,			()					)
     self.parent.assertEqual(result.start_top_log_probs.shape     ,			(self.batch_size, model.config.start_n_top)					)
     self.parent.assertEqual(result.start_top_index.shape     ,			(self.batch_size, model.config.start_n_top)					)
     self.parent.assertEqual(
         result.end_top_log_probs.shape     ,			(self.batch_size, model.config.start_n_top * model.config.end_n_top)					)
     self.parent.assertEqual(
         result.end_top_index.shape     ,			(self.batch_size, model.config.start_n_top * model.config.end_n_top)					)
     self.parent.assertEqual(result.cls_logits.shape     ,			(self.batch_size,)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			):
     __UpperCAmelCase					=   FlaubertForSequenceClassification(__A					)
     model.to(__A					)
     model.eval()

     __UpperCAmelCase					=   model(__A					)
     __UpperCAmelCase					=   model(__A     ,			labels=__A					)

     self.parent.assertEqual(result.loss.shape     ,			()					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.type_sequence_label_size)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			):
     __UpperCAmelCase					=   self.num_labels
     __UpperCAmelCase					=   FlaubertForTokenClassification(__A					)
     model.to(__A					)
     model.eval()

     __UpperCAmelCase					=   model(__A     ,			attention_mask=__A     ,			labels=__A					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.num_labels)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			):
     __UpperCAmelCase					=   self.num_choices
     __UpperCAmelCase					=   FlaubertForMultipleChoice(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   input_ids.unsqueeze(1					).expand(-1     ,			self.num_choices     ,			-1					).contiguous()
     __UpperCAmelCase					=   token_type_ids.unsqueeze(1					).expand(-1     ,			self.num_choices     ,			-1					).contiguous()
     __UpperCAmelCase					=   input_mask.unsqueeze(1					).expand(-1     ,			self.num_choices     ,			-1					).contiguous()
     __UpperCAmelCase					=   model(
         __A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			labels=__A     ,			)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_choices)					)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.prepare_config_and_inputs()
     (
         (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    
     )					=   config_and_inputs
     __UpperCAmelCase					=   {
         'input_ids': input_ids,
         'token_type_ids': token_type_ids,
         'lengths': input_lengths,
         'attention_mask': input_mask,
     }
     return config, inputs_dict






@require_torch
class UpperCAmelCase							(		UpperCAmelCase_					,			UpperCAmelCase_					,			unittest.TestCase							):
    _A :			List[str]						   =       (
        (
            FlaubertModel,
            FlaubertWithLMHeadModel,
            FlaubertForQuestionAnswering,
            FlaubertForQuestionAnsweringSimple,
            FlaubertForSequenceClassification,
            FlaubertForTokenClassification,
            FlaubertForMultipleChoice,
        )
        if is_torch_available()
        else ()
    )
    _A :			int						   =       (
        {
            """feature-extraction""": FlaubertModel,
            """fill-mask""": FlaubertWithLMHeadModel,
            """question-answering""": FlaubertForQuestionAnsweringSimple,
            """text-classification""": FlaubertForSequenceClassification,
            """token-classification""": FlaubertForTokenClassification,
            """zero-shot""": FlaubertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     if (
         pipeline_test_casse_name == "QAPipelineTests"
         and tokenizer_name is not None
         and not tokenizer_name.endswith('Fast'					)
     ):
      # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
      # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
      # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
      return True

     return False



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A=False					):
     __UpperCAmelCase					=   super()._prepare_for_class(__A     ,			__A     ,			return_labels=__A					)

     if return_labels:
      if model_class.__name__ == "FlaubertForQuestionAnswering":
       __UpperCAmelCase					=   torch.zeros(
           self.model_tester.batch_size     ,			dtype=torch.long     ,			device=__A					)
       __UpperCAmelCase					=   torch.zeros(
           self.model_tester.batch_size     ,			dtype=torch.long     ,			device=__A					)

     return inputs_dict



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   FlaubertModelTester(self					)
     __UpperCAmelCase					=   ConfigTester(self     ,			config_class=__A     ,			emb_dim=37					)



    def 					__lowerCamelCase						(  self					):
     self.config_tester.run_common_tests()



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_model(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_lm_head(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_simple_qa(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_qa(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_sequence_classif(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_token_classif(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_multiple_choice(*__A					)



    @slow
    def 					__lowerCamelCase						(  self					):
     for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      __UpperCAmelCase					=   FlaubertModel.from_pretrained(__A					)
      self.assertIsNotNone(__A					)




    @slow
    @require_torch_gpu
    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					,    __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes:
      # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
      if model_class == FlaubertForMultipleChoice:
       return

      __UpperCAmelCase					=   True
      __UpperCAmelCase					=   model_class(config=__A					)

      __UpperCAmelCase					=   self._prepare_for_class(__A     ,			__A					)
      __UpperCAmelCase					=   torch.jit.trace(
          __A     ,			(inputs_dict['input_ids'].to('cpu'					), inputs_dict['attention_mask'].to('cpu'					))					)

      with tempfile.TemporaryDirectory() as tmp:
       torch.jit.save(__A     ,			os.path.join(__A     ,			'traced_model.pt'					)					)
       __UpperCAmelCase					=   torch.jit.load(os.path.join(__A     ,			'traced_model.pt'					)     ,			map_location=__A					)
       loaded(inputs_dict['input_ids'].to(__A					)     ,			inputs_dict['attention_mask'].to(__A					)					)






@require_torch
class UpperCAmelCase							(		unittest.TestCase							):



    @slow
    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   FlaubertModel.from_pretrained('flaubert/flaubert_base_cased'					)
     __UpperCAmelCase					=   torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]					)
     with torch.no_grad():
      __UpperCAmelCase					=   model(__A					)[0]
     __UpperCAmelCase					=   torch.Size((1, 11, 768)					)
     self.assertEqual(output.shape     ,			__A					)
     __UpperCAmelCase					=   torch.tensor(
         [[[-2.6_2_5_1, -1.4_2_9_8, -0.0_2_2_7], [-2.8_5_1_0, -1.6_3_8_7, 0.2_2_5_8], [-2.8_1_1_4, -1.1_8_3_2, -0.3_0_6_6]]]					)

     self.assertTrue(torch.allclose(output[:, :3, :3]     ,			__A     ,			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


_A:					int							=				logging.getLogger(__name__)






class UpperCAmelCase							:



    def __init__(  self					):
     __UpperCAmelCase					=   False



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A					):
     if not self.initialized:
      __UpperCAmelCase					=   RagRetriever(
          __A     ,			question_encoder_tokenizer=__A     ,			generator_tokenizer=__A     ,			index=__A     ,			init_retrieval=__A     ,			)
      __UpperCAmelCase					=   True



    def 					__lowerCamelCase						(  self					):
     self.retriever.index.init_index()




    def 					__lowerCamelCase						(  self     ,			__A     ,			__A					):
     __UpperCAmelCase					,    __UpperCAmelCase					=   self.retriever._main_retrieve(__A     ,			__A					)
     return doc_ids, retrieved_doc_embeds






class UpperCAmelCase							(		UpperCAmelCase_							):



    def __init__(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A=None					):
     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					=   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 					__lowerCamelCase						(  self					):
     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 					__lowerCamelCase						(  self     ,			__A     ,			__A					):
     if len(self.retrieval_workers					) > 0:
      # Select a random retrieval actor.
      __UpperCAmelCase					=   self.retrieval_workers[random.randint(0     ,			len(self.retrieval_workers					) - 1					)]
      __UpperCAmelCase					,    __UpperCAmelCase					=   ray.get(random_worker.retrieve.remote(__A     ,			__A					)					)
     else:
      __UpperCAmelCase					,    __UpperCAmelCase					=   self._main_retrieve(__A     ,			__A					)
     return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(__A					)



    @classmethod
    def 					__lowerCamelCase						(  cls     ,			__A     ,			__A=None     ,			**__A					):
     return super(__A     ,			cls					).get_tokenizers(__A     ,			__A     ,			**__A					)




    @classmethod
    def 					__lowerCamelCase						(  cls     ,			__A     ,			__A     ,			__A=None     ,			**__A					):
     __UpperCAmelCase					=   kwargs.pop('config'     ,			__A					) or RagConfig.from_pretrained(__A     ,			**__A					)
     __UpperCAmelCase					=   RagTokenizer.from_pretrained(__A     ,			config=__A					)
     __UpperCAmelCase					=   rag_tokenizer.question_encoder
     __UpperCAmelCase					=   rag_tokenizer.generator
     if indexed_dataset is not None:
      __UpperCAmelCase					=   'custom'
      __UpperCAmelCase					=   CustomHFIndex(config.retrieval_vector_size     ,			__A					)
     else:
      __UpperCAmelCase					=   cls._build_index(__A					)
     return cls(
         __A     ,			question_encoder_tokenizer=__A     ,			generator_tokenizer=__A     ,			retrieval_workers=__A     ,			index=__A     ,			)





