





'''simple docstring'''


import unittest

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

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


if is_torch_available():
				import torch

				from transformers import (
				    RemBertConfig,
				    RemBertForCausalLM,
				    RemBertForMaskedLM,
				    RemBertForMultipleChoice,
				    RemBertForQuestionAnswering,
				    RemBertForSequenceClassification,
				    RemBertForTokenClassification,
				    RemBertModel,
				)
				from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST



class 							__UpperCAmelCase  :







		'''simple docstring'''







		def __init__(self						:      List[str]      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      Union[str, Any]=13      ,      _lowerCAmelCase						:      List[Any]=7      ,      _lowerCAmelCase						:      Tuple=True      ,      _lowerCAmelCase						:      str=True      ,      _lowerCAmelCase						:      Optional[int]=True      ,      _lowerCAmelCase						:      Tuple=True      ,      _lowerCAmelCase						:      Any=99      ,      _lowerCAmelCase						:      Optional[Any]=32      ,      _lowerCAmelCase						:      str=18      ,      _lowerCAmelCase						:      int=43      ,      _lowerCAmelCase						:      Optional[int]=5      ,      _lowerCAmelCase						:      List[str]=4      ,      _lowerCAmelCase						:      str=37      ,      _lowerCAmelCase						:      Optional[int]="gelu"      ,      _lowerCAmelCase						:      str=0.1      ,      _lowerCAmelCase						:      List[Any]=0.1      ,      _lowerCAmelCase						:      Union[str, Any]=512      ,      _lowerCAmelCase						:      Dict=16      ,      _lowerCAmelCase						:      Optional[int]=2      ,      _lowerCAmelCase						:      Dict=0.02      ,      _lowerCAmelCase						:      Any=3      ,      _lowerCAmelCase						:      Tuple=4      ,      _lowerCAmelCase						:      Any=None      ,      ):
			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							  =      input_embedding_size
			A							  =      output_embedding_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							  =      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							  =      ids_tensor([self.batch_size]      ,      self.num_choices      )

			A							  =      RemBertConfig(
			    vocab_size=self.vocab_size      ,      hidden_size=self.hidden_size      ,      input_embedding_size=self.input_embedding_size      ,      output_embedding_size=self.output_embedding_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      ,      is_decoder=_lowerCAmelCase      ,      initializer_range=self.initializer_range      ,      )

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







		def        A     (self						:      List[Any]      ):
			(
			    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    
			)							  =      self.prepare_config_and_inputs()

			A							  =      True
			A							  =      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
			A							  =      ids_tensor([self.batch_size, self.seq_length]      ,      vocab_size=2      )

			return (
			    config,
			    input_ids,
			    token_type_ids,
			    input_mask,
			    sequence_labels,
			    token_labels,
			    choice_labels,
			    encoder_hidden_states,
			    encoder_attention_mask,
			)







		def        A     (self						:      int      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Any      ):
			A							  =      RemBertModel(config=_lowerCAmelCase      )
			model.to(_lowerCAmelCase      )
			model.eval()
			A							  =      model(_lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      )
			A							  =      model(_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      )
			A							  =      model(_lowerCAmelCase      )
			self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)      )







		def        A     (self						:      List[str]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Union[str, Any]      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      int      ,      ):
			A							  =      True
			A							  =      RemBertModel(_lowerCAmelCase      )
			model.to(_lowerCAmelCase      )
			model.eval()
			A							  =      model(
			    _lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      encoder_hidden_states=_lowerCAmelCase      ,      encoder_attention_mask=_lowerCAmelCase      ,      )
			A							  =      model(
			    _lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      encoder_hidden_states=_lowerCAmelCase      ,      )
			A							  =      model(_lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      )
			self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)      )







		def        A     (self						:      List[Any]      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Union[str, Any]      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      Union[str, Any]      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Optional[int]      ,      ):
			A							  =      RemBertForCausalLM(config=_lowerCAmelCase      )
			model.to(_lowerCAmelCase      )
			model.eval()
			A							  =      model(_lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      labels=_lowerCAmelCase      )
			self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)      )







		def        A     (self						:      str      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      int      ):
			A							  =      RemBertForMaskedLM(config=_lowerCAmelCase      )
			model.to(_lowerCAmelCase      )
			model.eval()
			A							  =      model(_lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      labels=_lowerCAmelCase      )
			self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)      )







		def        A     (self						:      List[Any]      ,      _lowerCAmelCase						:      List[Any]      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Optional[Any]      ,      _lowerCAmelCase						:      List[str]      ,      ):
			A							  =      True
			A							  =      True
			A							  =      RemBertForCausalLM(config=_lowerCAmelCase      )
			model.to(_lowerCAmelCase      )
			model.eval()

			# first forward pass
			A							  =      model(
			    _lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      encoder_hidden_states=_lowerCAmelCase      ,      encoder_attention_mask=_lowerCAmelCase      ,      use_cache=_lowerCAmelCase      ,      )
			A							  =      outputs.past_key_values

			# create hypothetical multiple next token and extent to next_input_ids
			A							  =      ids_tensor((self.batch_size, 3)      ,      config.vocab_size      )
			A							  =      ids_tensor((self.batch_size, 3)      ,      vocab_size=2      )

			# append to next input_ids and
			A							  =      torch.cat([input_ids, next_tokens]      ,      dim=-1      )
			A							  =      torch.cat([input_mask, next_mask]      ,      dim=-1      )

			A							  =      model(
			    _lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      encoder_hidden_states=_lowerCAmelCase      ,      encoder_attention_mask=_lowerCAmelCase      ,      output_hidden_states=_lowerCAmelCase      ,      )["""hidden_states"""][0]
			A							  =      model(
			    _lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      encoder_hidden_states=_lowerCAmelCase      ,      encoder_attention_mask=_lowerCAmelCase      ,      past_key_values=_lowerCAmelCase      ,      output_hidden_states=_lowerCAmelCase      ,      )["""hidden_states"""][0]

			# select random slice
			A							  =      ids_tensor((1,)      ,      output_from_past.shape[-1]      ).item()
			A							  =      output_from_no_past[:, -3:, random_slice_idx].detach()
			A							  =      output_from_past[:, :, random_slice_idx].detach()

			self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]      )

			# test that outputs are equal for slice
			self.parent.assertTrue(torch.allclose(_lowerCAmelCase      ,      _lowerCAmelCase      ,      atol=1e-3      )      )







		def        A     (self						:      str      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      Union[str, Any]      ):
			A							  =      RemBertForQuestionAnswering(config=_lowerCAmelCase      )
			model.to(_lowerCAmelCase      )
			model.eval()
			A							  =      model(
			    _lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      start_positions=_lowerCAmelCase      ,      end_positions=_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        A     (self						:      Optional[int]      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      Any      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Dict      ):
			A							  =      self.num_labels
			A							  =      RemBertForSequenceClassification(_lowerCAmelCase      )
			model.to(_lowerCAmelCase      )
			model.eval()
			A							  =      model(_lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      labels=_lowerCAmelCase      )
			self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_labels)      )







		def        A     (self						:      int      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      str      ):
			A							  =      self.num_labels
			A							  =      RemBertForTokenClassification(config=_lowerCAmelCase      )
			model.to(_lowerCAmelCase      )
			model.eval()
			A							  =      model(_lowerCAmelCase      ,      attention_mask=_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      labels=_lowerCAmelCase      )
			self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.num_labels)      )







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







		def        A     (self						:      Optional[Any]      ):
			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_torch
class 							__UpperCAmelCase  (       A__		,					A__		,					unittest.TestCase						):







		'''simple docstring'''







		__lowerCAmelCase 				=    (
		    (
		        RemBertModel,
		        RemBertForMaskedLM,
		        RemBertForCausalLM,
		        RemBertForMultipleChoice,
		        RemBertForQuestionAnswering,
		        RemBertForSequenceClassification,
		        RemBertForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		__lowerCAmelCase 				=    (RemBertForCausalLM,) if is_torch_available() else ()
		__lowerCAmelCase 				=    (
		    {
		        '''feature-extraction''': RemBertModel,
		        '''fill-mask''': RemBertForMaskedLM,
		        '''question-answering''': RemBertForQuestionAnswering,
		        '''text-classification''': RemBertForSequenceClassification,
		        '''text-generation''': RemBertForCausalLM,
		        '''token-classification''': RemBertForTokenClassification,
		        '''zero-shot''': RemBertForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)







		def        A     (self						:      Tuple      ):
			A							  =      RemBertModelTester(self      )
			A							  =      ConfigTester(self      ,      config_class=_lowerCAmelCase      ,      hidden_size=37      )







		def        A     (self						:      Tuple      ):
			self.config_tester.run_common_tests()







		def        A     (self						:      Dict      ):
			A							  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*_lowerCAmelCase      )







		def        A     (self						:      List[Any]      ):
			A							  =      self.model_tester.prepare_config_and_inputs()
			for type in ["absolute", "relative_key", "relative_key_query"]:
				A							  =      type
				self.model_tester.create_and_check_model(*_lowerCAmelCase      )







		def        A     (self						:      int      ):
			A							  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase      )







		def        A     (self						:      List[str]      ):
			A							  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_multiple_choice(*_lowerCAmelCase      )







		def        A     (self						:      int      ):
			A							  =      self.model_tester.prepare_config_and_inputs_for_decoder()
			self.model_tester.create_and_check_decoder_model_past_large_inputs(*_lowerCAmelCase      )







		def        A     (self						:      Tuple      ):
			A							  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase      )







		def        A     (self						:      Any      ):
			A							  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_sequence_classification(*_lowerCAmelCase      )







		def        A     (self						:      Tuple      ):
			A							  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase      )







		def        A     (self						:      int      ):
			A							  =      self.model_tester.prepare_config_and_inputs_for_decoder()
			self.model_tester.create_and_check_model_as_decoder(*_lowerCAmelCase      )







		def        A     (self						:      str      ):
			# This regression test was failing with PyTorch < 1.3
			(
			    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    (
			    A
			)  ,    
			)							  =      self.model_tester.prepare_config_and_inputs_for_decoder()

			A							  =      None

			self.model_tester.create_and_check_model_as_decoder(
			    _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      )







		@slow
		def        A     (self						:      Optional[int]      ):
			for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
				A							  =      RemBertModel.from_pretrained(_lowerCAmelCase      )
				self.assertIsNotNone(_lowerCAmelCase      )






@require_torch
class 							__UpperCAmelCase  (       unittest.TestCase						):







		'''simple docstring'''







		@slow
		def        A     (self						:      List[str]      ):
			# Test exact values at the last hidden layer
			A							  =      RemBertModel.from_pretrained("""google/rembert"""      )
			A							  =      torch.tensor([[312, 5_6498, 313, 2125, 313]]      )
			A							  =      torch.tensor([[0, 0, 0, 1, 1]]      )
			with torch.no_grad():
				A							  =      model(_lowerCAmelCase      ,      token_type_ids=_lowerCAmelCase      ,      output_hidden_states=_lowerCAmelCase      )

			A							  =      1152

			A							  =      torch.Size((1, 5, hidden_size)      )
			self.assertEqual(output["""last_hidden_state"""].shape      ,      _lowerCAmelCase      )

			A							  =      torch.tensor(
			    [
			        [
			            [0.0_754, -0.2_022, 0.1_904],
			            [-0.3_354, -0.3_692, -0.4_791],
			            [-0.2_314, -0.6_729, -0.0_749],
			            [-0.0_396, -0.3_105, -0.4_234],
			            [-0.1_571, -0.0_525, 0.5_353],
			        ]
			    ]      )

			# Running on the original tf implementation gives slightly different results here.
			# Not clear why this variations is present
			# TODO: Find reason for discrepancy
			# expected_original_implementation = [[
			#     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
			#     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
			#     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
			#     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
			#     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
			# ]]

			self.assertTrue(torch.allclose(output["""last_hidden_state"""][:, :, :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				:	Tuple				    =       logging.getLogger(__name__)



class 							__UpperCAmelCase  :







		'''simple docstring'''







		def __init__(self						:      Union[str, Any]      ):
			A							  =      False







		def        A     (self						:      Any      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      str      ,      _lowerCAmelCase						:      Any      ):
			if not self.initialized:
				A							  =      RagRetriever(
				    _lowerCAmelCase      ,      question_encoder_tokenizer=_lowerCAmelCase      ,      generator_tokenizer=_lowerCAmelCase      ,      index=_lowerCAmelCase      ,      init_retrieval=_lowerCAmelCase      ,      )
				A							  =      True







		def        A     (self						:      Optional[int]      ):
			self.retriever.index.init_index()







		def        A     (self						:      Optional[int]      ):
			# delete the old self.retriever object before assigning the new index
			del self.retriever
			A							  =      False







		def        A     (self						:      List[str]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      List[str]      ):
			A  ,    A							  =      self.retriever._main_retrieve(_lowerCAmelCase      ,      _lowerCAmelCase      )
			A							  =      self.retriever.index.get_doc_dicts(_lowerCAmelCase      )
			return doc_ids, retrieved_doc_embeds, doc_dicts






class 							__UpperCAmelCase  (       A__						):







		'''simple docstring'''







		def __init__(self						:      Tuple      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      Tuple      ,      _lowerCAmelCase						:      str=None      ):
			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      ,      )

			A							  =      retrieval_workers
			A							  =      question_encoder_tokenizer
			A							  =      generator_tokenizer
			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						:      List[Any]      ):
			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						:      str      ,      _lowerCAmelCase						:      Dict      ,      _lowerCAmelCase						:      Tuple      ):
			if len(self.retrieval_workers      ) > 0:
				# Select a random retrieval actor.
				A							  =      self.retrieval_workers[random.randint(0      ,      len(self.retrieval_workers      ) - 1      )]
				A  ,    A  ,    A							  =      ray.get(
				    random_worker.retrieve.remote(_lowerCAmelCase      ,      _lowerCAmelCase      )      )
			else:
				A  ,    A							  =      self._main_retrieve(_lowerCAmelCase      ,      _lowerCAmelCase      )
				A							  =      self.index.get_doc_dicts(_lowerCAmelCase      )
			return retrieved_doc_embeds, doc_ids, doc_dicts







		@classmethod
		def        A     (cls						:      Union[str, Any]      ,      _lowerCAmelCase						:      int      ,      _lowerCAmelCase						:      Optional[Any]=None      ,      **_lowerCAmelCase						:      Union[str, Any]      ):
			return super(_lowerCAmelCase      ,      cls      ).get_tokenizers(_lowerCAmelCase      ,      _lowerCAmelCase      ,      **_lowerCAmelCase      )







		@classmethod
		def        A     (cls						:      List[Any]      ,      _lowerCAmelCase						:      Optional[int]      ,      _lowerCAmelCase						:      List[str]      ,      _lowerCAmelCase						:      str=None      ,      **_lowerCAmelCase						:      List[str]      ):
			A							  =      kwargs.pop("""config"""      ,      _lowerCAmelCase      ) or RagConfig.from_pretrained(_lowerCAmelCase      ,      **_lowerCAmelCase      )
			A							  =      RagTokenizer.from_pretrained(_lowerCAmelCase      ,      config=_lowerCAmelCase      )
			A							  =      rag_tokenizer.question_encoder
			A							  =      rag_tokenizer.generator

			if indexed_dataset is not None:
				A							  =      """custom"""
				A							  =      CustomHFIndex(config.retrieval_vector_size      ,      _lowerCAmelCase      )
			else:
				A							  =      cls._build_index(_lowerCAmelCase      )

			return cls(
			    _lowerCAmelCase      ,      question_encoder_tokenizer=_lowerCAmelCase      ,      generator_tokenizer=_lowerCAmelCase      ,      retrieval_workers=_lowerCAmelCase      ,      index=_lowerCAmelCase      ,      )







		def        A     (self						:      List[Any]      ):
			logger.info("""re-loading the new dataset with embeddings"""      )
			# access from the training loop

			ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]      )

			# build the index object again
			A							  =      self._build_index(self.config      )

			ray.get(
			    [
			        worker.create_rag_retriever.remote(
			            self.config      ,      self.question_encoder_tokenizer      ,      self.generator_tokenizer      ,      _lowerCAmelCase      )
			        for worker in self.retrieval_workers
			    ]      )

