



import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.generation import (
		    FlaxForcedBOSTokenLogitsProcessor,
		    FlaxForcedEOSTokenLogitsProcessor,
		    FlaxLogitsProcessorList,
		    FlaxMinLengthLogitsProcessor,
		    FlaxTemperatureLogitsWarper,
		    FlaxTopKLogitsWarper,
		    FlaxTopPLogitsWarper,
		)

@require_flax
class    snake_case__				(unittest.TestCase		):







						"""simple docstring"""




						def 							SCREAMING_SNAKE_CASE__(   self					,   __lowercase					,   __lowercase  )					->      Tuple:
									"""simple docstring"""




									a__   :						List[str]						       =  jnp.ones((batch_size, length)  ) / length
									return scores

						def 							SCREAMING_SNAKE_CASE__(   self  )					->      Union[str, Any]:
									"""simple docstring"""




									a__   :						Tuple						       =  None
									a__   :						str						       =  2_0

									a__   :						Tuple						       =  self._get_uniform_logits(batch_size=2					,   length=__SCREAMING_SNAKE_CASE  )

									# tweak scores to not be uniform anymore
									a__   :						List[str]						       =  scores.at[1, 5].set((1 / length) + 0.1  )  # peak, 1st batch
									a__   :						Any						       =  scores.at[1, 1_0].set((1 / length) - 0.4  )  # valley, 1st batch

									# compute softmax
									a__   :						str						       =  jax.nn.softmax(__SCREAMING_SNAKE_CASE					,   axis=-1  )

									a__   :						List[str]						       =  FlaxTemperatureLogitsWarper(temperature=0.5  )
									a__   :						List[Any]						       =  FlaxTemperatureLogitsWarper(temperature=1.3  )

									a__   :						Tuple						       =  jax.nn.softmax(temp_dist_warper_sharper(__SCREAMING_SNAKE_CASE					,   scores.copy()					,   cur_len=__SCREAMING_SNAKE_CASE  )					,   axis=-1  )
									a__   :						Optional[Any]						       =  jax.nn.softmax(temp_dist_warper_smoother(__SCREAMING_SNAKE_CASE					,   scores.copy()					,   cur_len=__SCREAMING_SNAKE_CASE  )					,   axis=-1  )

									# uniform distribution stays uniform
									self.assertTrue(jnp.allclose(probs[0, :]					,   warped_prob_sharp[0, :]					,   atol=1E-3  )  )
									self.assertTrue(jnp.allclose(probs[0, :]					,   warped_prob_smooth[0, :]					,   atol=1E-3  )  )

									# sharp peaks get higher, valleys get lower
									self.assertLess(probs[1, :].max()					,   warped_prob_sharp[1, :].max()  )
									self.assertGreater(probs[1, :].min()					,   warped_prob_sharp[1, :].min()  )

									# smooth peaks get lower, valleys get higher
									self.assertGreater(probs[1, :].max()					,   warped_prob_smooth[1, :].max()  )
									self.assertLess(probs[1, :].min()					,   warped_prob_smooth[1, :].min()  )

						def 							SCREAMING_SNAKE_CASE__(   self  )					->      List[Any]:
									"""simple docstring"""




									a__   :						Optional[Any]						       =  None
									a__   :						Dict						       =  1_0
									a__   :						List[Any]						       =  2

									# create ramp distribution
									a__   :						Any						       =  np.broadcast_to(np.arange(__SCREAMING_SNAKE_CASE  )[None, :]					,   (batch_size, vocab_size)  ).copy()
									a__   :						Optional[Any]						       =  ramp_logits[1:, : vocab_size // 2] + vocab_size

									a__   :						int						       =  FlaxTopKLogitsWarper(3  )

									a__   :						Any						       =  top_k_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )

									# check that correct tokens are filtered
									self.assertListEqual(jnp.isinf(scores[0]  ).tolist()					,   7 * [True] + 3 * [False]  )
									self.assertListEqual(jnp.isinf(scores[1]  ).tolist()					,   2 * [True] + 3 * [False] + 5 * [True]  )

									# check special case
									a__   :						int						       =  5
									a__   :						Union[str, Any]						       =  FlaxTopKLogitsWarper(top_k=1					,   filter_value=0.0					,   min_tokens_to_keep=3  )

									a__   :						List[Any]						       =  np.broadcast_to(np.arange(__SCREAMING_SNAKE_CASE  )[None, :]					,   (batch_size, length)  ).copy()
									a__   :						List[str]						       =  top_k_warp_safety_check(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )

									# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
									self.assertListEqual((scores == 0.0).sum(axis=-1  ).tolist()					,   [2, 2]  )

						def 							SCREAMING_SNAKE_CASE__(   self  )					->      List[Any]:
									"""simple docstring"""




									a__   :						Optional[int]						       =  None
									a__   :						Any						       =  1_0
									a__   :						Any						       =  2

									# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
									a__   :						Any						       =  np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.1_5, 0.3, 0.3, 0.2_5]]  )  )

									a__   :						Optional[Any]						       =  FlaxTopPLogitsWarper(0.8  )
									a__   :						Optional[Any]						       =  np.exp(top_p_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )  )

									# dist should be filtered to keep min num values so that sum is >= top_p
									# exp (-inf) => 0
									a__   :						Optional[int]						       =  np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]  )
									self.assertTrue(np.allclose(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   atol=1E-3  )  )

									# check edge cases with negative and extreme logits
									a__   :						str						       =  np.broadcast_to(np.arange(__SCREAMING_SNAKE_CASE  )[None, :]					,   (batch_size, vocab_size)  ).copy() - (
									    vocab_size // 2
									)

									# make ramp_logits more extreme
									a__   :						Optional[int]						       =  ramp_logits[1] * 1_0_0.0

									# make sure at least 2 tokens are kept
									a__   :						List[str]						       =  FlaxTopPLogitsWarper(0.9					,   min_tokens_to_keep=2					,   filter_value=0.0  )
									a__   :						Union[str, Any]						       =  top_p_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )

									# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
									self.assertListEqual((filtered_dist != 0.0).sum(axis=-1  ).tolist()					,   [3, 2]  )

						def 							SCREAMING_SNAKE_CASE__(   self  )					->      Optional[int]:
									"""simple docstring"""




									a__   :						int						       =  2_0
									a__   :						List[Any]						       =  4
									a__   :						List[Any]						       =  0

									a__   :						Optional[int]						       =  FlaxMinLengthLogitsProcessor(min_length=1_0					,   eos_token_id=__SCREAMING_SNAKE_CASE  )

									# check that min length is applied at length 5
									a__   :						Any						       =  ids_tensor((batch_size, 2_0)					,   vocab_size=2_0  )
									a__   :						List[Any]						       =  5
									a__   :						Any						       =  self._get_uniform_logits(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						List[str]						       =  min_dist_processor(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()					,   4 * [-float("""inf"""  )]  )

									# check that min length is not applied anymore at length 15
									a__   :						Optional[int]						       =  self._get_uniform_logits(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						Any						       =  1_5
									a__   :						Any						       =  min_dist_processor(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									self.assertFalse(jnp.isinf(__SCREAMING_SNAKE_CASE  ).any()  )

						def 							SCREAMING_SNAKE_CASE__(   self  )					->      List[Any]:
									"""simple docstring"""




									a__   :						Union[str, Any]						       =  2_0
									a__   :						List[Any]						       =  4
									a__   :						Optional[Any]						       =  0

									a__   :						Optional[Any]						       =  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE  )

									# check that all scores are -inf except the bos_token_id score
									a__   :						Optional[Any]						       =  ids_tensor((batch_size, 1)					,   vocab_size=2_0  )
									a__   :						List[Any]						       =  1
									a__   :						Optional[int]						       =  self._get_uniform_logits(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						Dict						       =  logits_processor(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]  ).all()  )
									self.assertListEqual(scores[:, bos_token_id].tolist()					,   4 * [0]  )  # score for bos_token_id shold be zero

									# check that bos_token_id is not forced if current length is greater than 1
									a__   :						List[str]						       =  3
									a__   :						Dict						       =  self._get_uniform_logits(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						int						       =  logits_processor(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									self.assertFalse(jnp.isinf(__SCREAMING_SNAKE_CASE  ).any()  )

						def 							SCREAMING_SNAKE_CASE__(   self  )					->      List[str]:
									"""simple docstring"""




									a__   :						Optional[int]						       =  2_0
									a__   :						Union[str, Any]						       =  4
									a__   :						str						       =  0
									a__   :						Tuple						       =  5

									a__   :						List[Any]						       =  FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE					,   eos_token_id=__SCREAMING_SNAKE_CASE  )

									# check that all scores are -inf except the eos_token_id when max_length is reached
									a__   :						List[str]						       =  ids_tensor((batch_size, 4)					,   vocab_size=2_0  )
									a__   :						Optional[Any]						       =  4
									a__   :						List[Any]						       =  self._get_uniform_logits(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						Any						       =  logits_processor(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]  ).all()  )
									self.assertListEqual(scores[:, eos_token_id].tolist()					,   4 * [0]  )  # score for eos_token_id should be zero

									# check that eos_token_id is not forced if max_length is not reached
									a__   :						List[str]						       =  3
									a__   :						str						       =  self._get_uniform_logits(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						int						       =  logits_processor(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									self.assertFalse(jnp.isinf(__SCREAMING_SNAKE_CASE  ).any()  )

						def 							SCREAMING_SNAKE_CASE__(   self  )					->      str:
									"""simple docstring"""




									a__   :						Optional[Any]						       =  4
									a__   :						int						       =  1_0
									a__   :						Dict						       =  1_5
									a__   :						List[Any]						       =  2
									a__   :						int						       =  1
									a__   :						Optional[int]						       =  1_5

									# dummy input_ids and scores
									a__   :						Optional[int]						       =  ids_tensor((batch_size, sequence_length)					,   __SCREAMING_SNAKE_CASE  )
									a__   :						Tuple						       =  input_ids.copy()

									a__   :						Union[str, Any]						       =  self._get_uniform_logits(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						Any						       =  scores.copy()

									# instantiate all dist processors
									a__   :						List[str]						       =  FlaxTemperatureLogitsWarper(temperature=0.5  )
									a__   :						List[str]						       =  FlaxTopKLogitsWarper(3  )
									a__   :						Union[str, Any]						       =  FlaxTopPLogitsWarper(0.8  )

									# instantiate all logits processors
									a__   :						Tuple						       =  FlaxMinLengthLogitsProcessor(min_length=1_0					,   eos_token_id=__SCREAMING_SNAKE_CASE  )
									a__   :						Tuple						       =  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE  )
									a__   :						int						       =  FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE					,   eos_token_id=__SCREAMING_SNAKE_CASE  )

									a__   :						List[str]						       =  1_0

									# no processor list
									a__   :						Dict						       =  temp_dist_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									a__   :						Optional[int]						       =  top_k_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									a__   :						List[Any]						       =  top_p_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									a__   :						Optional[Any]						       =  min_dist_proc(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									a__   :						Tuple						       =  bos_dist_proc(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
									a__   :						Tuple						       =  eos_dist_proc(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )

									# with processor list
									a__   :						Union[str, Any]						       =  FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
									a__   :						Tuple						       =  processor(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )

									# scores should be equal
									self.assertTrue(jnp.allclose(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   atol=1E-3  )  )

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist()					,   input_ids_comp.tolist()  )

						def 							SCREAMING_SNAKE_CASE__(   self  )					->      Tuple:
									"""simple docstring"""




									a__   :						Optional[Any]						       =  4
									a__   :						Union[str, Any]						       =  1_0
									a__   :						Union[str, Any]						       =  1_5
									a__   :						Any						       =  2
									a__   :						Optional[Any]						       =  1
									a__   :						Tuple						       =  1_5

									# dummy input_ids and scores
									a__   :						int						       =  ids_tensor((batch_size, sequence_length)					,   __SCREAMING_SNAKE_CASE  )
									a__   :						Union[str, Any]						       =  input_ids.copy()

									a__   :						Optional[int]						       =  self._get_uniform_logits(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						List[str]						       =  scores.copy()

									# instantiate all dist processors
									a__   :						Optional[Any]						       =  FlaxTemperatureLogitsWarper(temperature=0.5  )
									a__   :						Optional[int]						       =  FlaxTopKLogitsWarper(3  )
									a__   :						List[str]						       =  FlaxTopPLogitsWarper(0.8  )

									# instantiate all logits processors
									a__   :						int						       =  FlaxMinLengthLogitsProcessor(min_length=1_0					,   eos_token_id=__SCREAMING_SNAKE_CASE  )
									a__   :						int						       =  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE  )
									a__   :						Union[str, Any]						       =  FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE					,   eos_token_id=__SCREAMING_SNAKE_CASE  )

									a__   :						Optional[Any]						       =  1_0

									# no processor list
									def run_no_processor_list(__lowercase					,   __lowercase					,   __lowercase  ):
												a__   :						List[Any]						       =  temp_dist_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
												a__   :						Dict						       =  top_k_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
												a__   :						Optional[Any]						       =  top_p_warp(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
												a__   :						Dict						       =  min_dist_proc(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
												a__   :						List[str]						       =  bos_dist_proc(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
												a__   :						List[str]						       =  eos_dist_proc(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
												return scores

									# with processor list
									def run_processor_list(__lowercase					,   __lowercase					,   __lowercase  ):
												a__   :						List[str]						       =  FlaxLogitsProcessorList(
												    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
												a__   :						Union[str, Any]						       =  processor(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   cur_len=__SCREAMING_SNAKE_CASE  )
												return scores

									a__   :						Dict						       =  jax.jit(__SCREAMING_SNAKE_CASE  )
									a__   :						Dict						       =  jax.jit(__SCREAMING_SNAKE_CASE  )

									a__   :						List[Any]						       =  jitted_run_no_processor_list(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )
									a__   :						Union[str, Any]						       =  jitted_run_processor_list(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE  )

									# scores should be equal
									self.assertTrue(jnp.allclose(__SCREAMING_SNAKE_CASE					,   __SCREAMING_SNAKE_CASE					,   atol=1E-3  )  )

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist()					,   input_ids_comp.tolist()  )









import logging
import random

import ray

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


__snake_case						:List[Any]      		=     logging.getLogger(__name__)

class     _A       :



       def __init__(		self							:     List[str]):


           '''simple docstring'''

           __a						       =				False



       def 		_lowerCamelCase			(		self							:     Any     ,		__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,		__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,		__SCREAMING_SNAKE_CASE							:     List[str]     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

           self.retriever.index.init_index()





       def 		_lowerCamelCase			(		self							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]):


           '''simple docstring'''

           __a       ,   __a						       =				self.retriever._main_retrieve(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           return doc_ids, retrieved_doc_embeds


class     _A       (	__UpperCAmelCase    ):



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


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''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 		_lowerCamelCase			(		self							:     Dict     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     List[Any]):


           '''simple docstring'''

           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						       =				ray.get(random_worker.retrieve.remote(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE))
           else:
               __a       ,   __a						       =				self._main_retrieve(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(__SCREAMING_SNAKE_CASE)



       @classmethod
       def 		_lowerCamelCase			(		cls							:     Any     ,		__SCREAMING_SNAKE_CASE							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     Tuple=None     ,		**__SCREAMING_SNAKE_CASE							:     Optional[int]):


           '''simple docstring'''

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





       @classmethod
       def 		_lowerCamelCase			(		cls							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     str     ,		__SCREAMING_SNAKE_CASE							:     str=None     ,		**__SCREAMING_SNAKE_CASE							:     List[Any]):


           '''simple docstring'''

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





