
"""simple docstring"""



import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





class 				a			(					lowercase						,		unittest.TestCase      ):
      UpperCamelCase							:   Union[str, Any]           =     AudioLDMPipeline
      UpperCamelCase							:   Optional[int]           =     TEXT_TO_AUDIO_PARAMS
      UpperCamelCase							:   List[str]           =     TEXT_TO_AUDIO_BATCH_PARAMS
      UpperCamelCase							:   List[str]           =     frozenset(
          [
              """num_inference_steps""",
              """num_waveforms_per_prompt""",
              """generator""",
              """latents""",
              """output_type""",
              """return_dict""",
              """callback""",
              """callback_steps""",
          ]      )



      def  __snake_case							(   self      ):
          torch.manual_seed(0      )
          UpperCAmelCase__	:      Union[str, Any]     							=					UNetaDConditionModel(
              block_out_channels=(32, 64)     ,	layers_per_block=2     ,	sample_size=32     ,	in_channels=4     ,	out_channels=4     ,	down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')     ,	up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')     ,	cross_attention_dim=(32, 64)     ,	class_embed_type='simple_projection'     ,	projection_class_embeddings_input_dim=32     ,	class_embeddings_concat=UpperCamelCase_     ,	)
          UpperCAmelCase__	:      Union[str, Any]     							=					DDIMScheduler(
              beta_start=0.00085     ,	beta_end=0.012     ,	beta_schedule='scaled_linear'     ,	clip_sample=UpperCamelCase_     ,	set_alpha_to_one=UpperCamelCase_     ,	)
          torch.manual_seed(0      )
          UpperCAmelCase__	:      List[str]     							=					AutoencoderKL(
              block_out_channels=[32, 64]     ,	in_channels=1     ,	out_channels=1     ,	down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']     ,	up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']     ,	latent_channels=4     ,	)
          torch.manual_seed(0      )
          UpperCAmelCase__	:      Optional[Any]     							=					ClapTextConfig(
              bos_token_id=0     ,	eos_token_id=2     ,	hidden_size=32     ,	intermediate_size=37     ,	layer_norm_eps=1E-05     ,	num_attention_heads=4     ,	num_hidden_layers=5     ,	pad_token_id=1     ,	vocab_size=1_000     ,	projection_dim=32     ,	)
          UpperCAmelCase__	:      Any     							=					ClapTextModelWithProjection(UpperCamelCase_      )
          UpperCAmelCase__	:      List[Any]     							=					RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'     ,	model_max_length=77      )

          UpperCAmelCase__	:      Any     							=					SpeechTaHifiGanConfig(
              model_in_dim=8     ,	sampling_rate=16_000     ,	upsample_initial_channel=16     ,	upsample_rates=[2, 2]     ,	upsample_kernel_sizes=[4, 4]     ,	resblock_kernel_sizes=[3, 7]     ,	resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]     ,	normalize_before=UpperCamelCase_     ,	)

          UpperCAmelCase__	:      Dict     							=					SpeechTaHifiGan(UpperCamelCase_      )

          UpperCAmelCase__	:      List[str]     							=					{
              'unet': unet,
              'scheduler': scheduler,
              'vae': vae,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'vocoder': vocoder,
          }
          return components



      def  __snake_case							(   self     ,	UpperCamelCase_     ,	UpperCamelCase_=0      ):
          if str(UpperCamelCase_      ).startswith('mps'      ):
              UpperCAmelCase__	:      Dict     							=					torch.manual_seed(UpperCamelCase_      )
          else:
              UpperCAmelCase__	:      Union[str, Any]     							=					torch.Generator(device=UpperCamelCase_      ).manual_seed(UpperCamelCase_      )
          UpperCAmelCase__	:      List[str]     							=					{
              'prompt': 'A hammer hitting a wooden surface',
              'generator': generator,
              'num_inference_steps': 2,
              'guidance_scale': 6.0,
          }
          return inputs



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Any     							=					'cpu'  # ensure determinism for the device-dependent torch.Generator

          UpperCAmelCase__	:      Optional[Any]     							=					self.get_dummy_components()
          UpperCAmelCase__	:      Optional[int]     							=					AudioLDMPipeline(**UpperCamelCase_      )
          UpperCAmelCase__	:      List[str]     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )

          UpperCAmelCase__	:      Optional[Any]     							=					self.get_dummy_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      int     							=					audioldm_pipe(**UpperCamelCase_      )
          UpperCAmelCase__	:      int     							=					output.audios[0]

          assert audio.ndim == 1
          assert len(UpperCamelCase_      ) == 256

          UpperCAmelCase__	:      List[str]     							=					audio[:10]
          UpperCAmelCase__	:      List[str]     							=					np.array(
              [-0.0050, 0.0050, -0.0060, 0.0033, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0033]      )

          assert np.abs(audio_slice - expected_slice      ).max() < 1E-2



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      str     							=					self.get_dummy_components()
          UpperCAmelCase__	:      List[str]     							=					AudioLDMPipeline(**UpperCamelCase_      )
          UpperCAmelCase__	:      Any     							=					audioldm_pipe.to(UpperCamelCase_      )
          UpperCAmelCase__	:      List[Any]     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )

          UpperCAmelCase__	:      Any     							=					self.get_dummy_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      List[str]     							=					3 * [inputs['prompt']]

          # forward
          UpperCAmelCase__	:      str     							=					audioldm_pipe(**UpperCamelCase_      )
          UpperCAmelCase__	:      str     							=					output.audios[0]

          UpperCAmelCase__	:      List[Any]     							=					self.get_dummy_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      List[str]     							=					3 * [inputs.pop('prompt'      )]

          UpperCAmelCase__	:      Union[str, Any]     							=					audioldm_pipe.tokenizer(
              UpperCamelCase_     ,	padding='max_length'     ,	max_length=audioldm_pipe.tokenizer.model_max_length     ,	truncation=UpperCamelCase_     ,	return_tensors='pt'     ,	)
          UpperCAmelCase__	:      Optional[int]     							=					text_inputs['input_ids'].to(UpperCamelCase_      )

          UpperCAmelCase__	:      List[str]     							=					audioldm_pipe.text_encoder(
              UpperCamelCase_     ,	)
          UpperCAmelCase__	:      List[str]     							=					prompt_embeds.text_embeds
          # additional L_2 normalization over each hidden-state
          UpperCAmelCase__	:      Dict     							=					F.normalize(UpperCamelCase_     ,	dim=-1      )

          UpperCAmelCase__	:      str     							=					prompt_embeds

          # forward
          UpperCAmelCase__	:      Dict     							=					audioldm_pipe(**UpperCamelCase_      )
          UpperCAmelCase__	:      str     							=					output.audios[0]

          assert np.abs(audio_a - audio_a      ).max() < 1E-2



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Dict     							=					self.get_dummy_components()
          UpperCAmelCase__	:      int     							=					AudioLDMPipeline(**UpperCamelCase_      )
          UpperCAmelCase__	:      Optional[int]     							=					audioldm_pipe.to(UpperCamelCase_      )
          UpperCAmelCase__	:      str     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )

          UpperCAmelCase__	:      Tuple     							=					self.get_dummy_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      Optional[Any]     							=					3 * ['this is a negative prompt']
          UpperCAmelCase__	:      Any     							=					negative_prompt
          UpperCAmelCase__	:      Any     							=					3 * [inputs['prompt']]

          # forward
          UpperCAmelCase__	:      Optional[int]     							=					audioldm_pipe(**UpperCamelCase_      )
          UpperCAmelCase__	:      int     							=					output.audios[0]

          UpperCAmelCase__	:      str     							=					self.get_dummy_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      Optional[int]     							=					3 * [inputs.pop('prompt'      )]

          UpperCAmelCase__	:      str     							=					[]
          for p in [prompt, negative_prompt]:
              UpperCAmelCase__	:      Optional[int]     							=					audioldm_pipe.tokenizer(
                  UpperCamelCase_     ,	padding='max_length'     ,	max_length=audioldm_pipe.tokenizer.model_max_length     ,	truncation=UpperCamelCase_     ,	return_tensors='pt'     ,	)
              UpperCAmelCase__	:      Dict     							=					text_inputs['input_ids'].to(UpperCamelCase_      )

              UpperCAmelCase__	:      Dict     							=					audioldm_pipe.text_encoder(
                  UpperCamelCase_     ,	)
              UpperCAmelCase__	:      Optional[int]     							=					text_embeds.text_embeds
              # additional L_2 normalization over each hidden-state
              UpperCAmelCase__	:      Any     							=					F.normalize(UpperCamelCase_     ,	dim=-1      )

              embeds.append(UpperCamelCase_      )

          UpperCAmelCase__							,							UpperCAmelCase__	:      List[str]     							=					embeds

          # forward
          UpperCAmelCase__	:      List[Any]     							=					audioldm_pipe(**UpperCamelCase_      )
          UpperCAmelCase__	:      int     							=					output.audios[0]

          assert np.abs(audio_a - audio_a      ).max() < 1E-2



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      List[str]     							=					'cpu'  # ensure determinism for the device-dependent torch.Generator
          UpperCAmelCase__	:      str     							=					self.get_dummy_components()
          UpperCAmelCase__	:      int     							=					PNDMScheduler(skip_prk_steps=UpperCamelCase_      )
          UpperCAmelCase__	:      Optional[Any]     							=					AudioLDMPipeline(**UpperCamelCase_      )
          UpperCAmelCase__	:      Union[str, Any]     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )

          UpperCAmelCase__	:      str     							=					self.get_dummy_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      str     							=					'egg cracking'
          UpperCAmelCase__	:      List[str]     							=					audioldm_pipe(**UpperCamelCase_     ,	negative_prompt=UpperCamelCase_      )
          UpperCAmelCase__	:      List[str]     							=					output.audios[0]

          assert audio.ndim == 1
          assert len(UpperCamelCase_      ) == 256

          UpperCAmelCase__	:      List[Any]     							=					audio[:10]
          UpperCAmelCase__	:      List[Any]     							=					np.array(
              [-0.0051, 0.0050, -0.0060, 0.0034, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0032]      )

          assert np.abs(audio_slice - expected_slice      ).max() < 1E-2



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Union[str, Any]     							=					'cpu'  # ensure determinism for the device-dependent torch.Generator
          UpperCAmelCase__	:      Tuple     							=					self.get_dummy_components()
          UpperCAmelCase__	:      Tuple     							=					PNDMScheduler(skip_prk_steps=UpperCamelCase_      )
          UpperCAmelCase__	:      Any     							=					AudioLDMPipeline(**UpperCamelCase_      )
          UpperCAmelCase__	:      Tuple     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )

          UpperCAmelCase__	:      Dict     							=					'A hammer hitting a wooden surface'

          # test num_waveforms_per_prompt=1 (default)
          UpperCAmelCase__	:      Optional[Any]     							=					audioldm_pipe(UpperCamelCase_     ,	num_inference_steps=2      ).audios

          assert audios.shape == (1, 256)

          # test num_waveforms_per_prompt=1 (default) for batch of prompts
          UpperCAmelCase__	:      Dict     							=					2
          UpperCAmelCase__	:      Optional[int]     							=					audioldm_pipe([prompt] * batch_size     ,	num_inference_steps=2      ).audios

          assert audios.shape == (batch_size, 256)

          # test num_waveforms_per_prompt for single prompt
          UpperCAmelCase__	:      Tuple     							=					2
          UpperCAmelCase__	:      Dict     							=					audioldm_pipe(UpperCamelCase_     ,	num_inference_steps=2     ,	num_waveforms_per_prompt=UpperCamelCase_      ).audios

          assert audios.shape == (num_waveforms_per_prompt, 256)

          # test num_waveforms_per_prompt for batch of prompts
          UpperCAmelCase__	:      List[Any]     							=					2
          UpperCAmelCase__	:      Dict     							=					audioldm_pipe(
              [prompt] * batch_size     ,	num_inference_steps=2     ,	num_waveforms_per_prompt=UpperCamelCase_      ).audios

          assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Optional[Any]     							=					'cpu'  # ensure determinism for the device-dependent torch.Generator
          UpperCAmelCase__	:      Union[str, Any]     							=					self.get_dummy_components()
          UpperCAmelCase__	:      Any     							=					AudioLDMPipeline(**UpperCamelCase_      )
          UpperCAmelCase__	:      Union[str, Any]     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )
          UpperCAmelCase__	:      int     							=					audioldm_pipe.vocoder.config.sampling_rate

          UpperCAmelCase__	:      List[Any]     							=					self.get_dummy_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      Optional[Any]     							=					audioldm_pipe(audio_length_in_s=0.016     ,	**UpperCamelCase_      )
          UpperCAmelCase__	:      Optional[int]     							=					output.audios[0]

          assert audio.ndim == 1
          assert len(UpperCamelCase_      ) / vocoder_sampling_rate == 0.016

          UpperCAmelCase__	:      int     							=					audioldm_pipe(audio_length_in_s=0.032     ,	**UpperCamelCase_      )
          UpperCAmelCase__	:      List[str]     							=					output.audios[0]

          assert audio.ndim == 1
          assert len(UpperCamelCase_      ) / vocoder_sampling_rate == 0.032



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Optional[Any]     							=					self.get_dummy_components()
          UpperCAmelCase__	:      Optional[Any]     							=					AudioLDMPipeline(**UpperCamelCase_      )
          UpperCAmelCase__	:      Optional[int]     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )

          UpperCAmelCase__	:      Optional[Any]     							=					['hey']

          UpperCAmelCase__	:      Union[str, Any]     							=					audioldm_pipe(UpperCamelCase_     ,	num_inference_steps=1      )
          UpperCAmelCase__	:      Union[str, Any]     							=					output.audios.shape
          assert audio_shape == (1, 256)

          UpperCAmelCase__	:      str     							=					audioldm_pipe.vocoder.config
          config.model_in_dim *= 2
          UpperCAmelCase__	:      str     							=					SpeechTaHifiGan(UpperCamelCase_      ).to(UpperCamelCase_      )
          UpperCAmelCase__	:      List[Any]     							=					audioldm_pipe(UpperCamelCase_     ,	num_inference_steps=1      )
          UpperCAmelCase__	:      Optional[Any]     							=					output.audios.shape
          # waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
          assert audio_shape == (1, 256)



      def  __snake_case							(   self      ):
          self._test_attention_slicing_forward_pass(test_mean_pixel_difference=UpperCamelCase_      )



      def  __snake_case							(   self      ):
          self._test_inference_batch_single_identical(test_mean_pixel_difference=UpperCamelCase_      )



      @unittest.skipIf(
          torch_device != 'cuda' or not is_xformers_available()     ,	reason='XFormers attention is only available with CUDA and `xformers` installed'     ,	)
      def  __snake_case							(   self      ):
          self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=UpperCamelCase_      )





@slow
class 				a			(					unittest.TestCase      ):



      def  __snake_case							(   self      ):
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



      def  __snake_case							(   self     ,	UpperCamelCase_     ,	UpperCamelCase_="cpu"     ,	UpperCamelCase_=torch.floataa     ,	UpperCamelCase_=0      ):
          UpperCAmelCase__	:      int     							=					torch.Generator(device=UpperCamelCase_      ).manual_seed(UpperCamelCase_      )
          UpperCAmelCase__	:      Optional[int]     							=					np.random.RandomState(UpperCamelCase_      ).standard_normal((1, 8, 128, 16)      )
          UpperCAmelCase__	:      Union[str, Any]     							=					torch.from_numpy(UpperCamelCase_      ).to(device=UpperCamelCase_     ,	dtype=UpperCamelCase_      )
          UpperCAmelCase__	:      List[Any]     							=					{
              'prompt': 'A hammer hitting a wooden surface',
              'latents': latents,
              'generator': generator,
              'num_inference_steps': 3,
              'guidance_scale': 2.5,
          }
          return inputs



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      List[str]     							=					AudioLDMPipeline.from_pretrained('cvssp/audioldm'      )
          UpperCAmelCase__	:      int     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )

          UpperCAmelCase__	:      Tuple     							=					self.get_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      str     							=					25
          UpperCAmelCase__	:      int     							=					audioldm_pipe(**UpperCamelCase_      ).audios[0]

          assert audio.ndim == 1
          assert len(UpperCamelCase_      ) == 81_920

          UpperCAmelCase__	:      Optional[Any]     							=					audio[77_230:77_240]
          UpperCAmelCase__	:      Union[str, Any]     							=					np.array(
              [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315]      )
          UpperCAmelCase__	:      Tuple     							=					np.abs(expected_slice - audio_slice      ).max()
          assert max_diff < 1E-2



      def  __snake_case							(   self      ):
          UpperCAmelCase__	:      Any     							=					AudioLDMPipeline.from_pretrained('cvssp/audioldm'      )
          UpperCAmelCase__	:      List[Any]     							=					LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config      )
          UpperCAmelCase__	:      List[Any]     							=					audioldm_pipe.to(UpperCamelCase_      )
          audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase_      )

          UpperCAmelCase__	:      Any     							=					self.get_inputs(UpperCamelCase_      )
          UpperCAmelCase__	:      Dict     							=					audioldm_pipe(**UpperCamelCase_      ).audios[0]

          assert audio.ndim == 1
          assert len(UpperCamelCase_      ) == 81_920

          UpperCAmelCase__	:      Optional[int]     							=					audio[27_780:27_790]
          UpperCAmelCase__	:      Optional[Any]     							=					np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212]      )
          UpperCAmelCase__	:      Tuple     							=					np.abs(expected_slice - audio_slice      ).max()
          assert max_diff < 3E-2



"""simple docstring"""



import logging
import random

import ray

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


UpperCamelCase__          =	logging.getLogger(__name__)





class 				a			:



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



      def  __snake_case							(   self     ,	UpperCamelCase_     ,	UpperCamelCase_     ,	UpperCamelCase_     ,	UpperCamelCase_      ):
          if not self.initialized:
              UpperCAmelCase__	:      Optional[Any]     							=					RagRetriever(
                  UpperCamelCase_     ,	question_encoder_tokenizer=UpperCamelCase_     ,	generator_tokenizer=UpperCamelCase_     ,	index=UpperCamelCase_     ,	init_retrieval=UpperCamelCase_     ,	)
              UpperCAmelCase__	:      Union[str, Any]     							=					True



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



      def  __snake_case							(   self     ,	UpperCamelCase_     ,	UpperCamelCase_      ):
          UpperCAmelCase__							,							UpperCAmelCase__	:      Union[str, Any]     							=					self.retriever._main_retrieve(UpperCamelCase_     ,	UpperCamelCase_      )
          return doc_ids, retrieved_doc_embeds





class 				a			(					lowercase      ):



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



      def  __snake_case							(   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  __snake_case							(   self     ,	UpperCamelCase_     ,	UpperCamelCase_      ):
          if len(self.retrieval_workers      ) > 0:
              # Select a random retrieval actor.
              UpperCAmelCase__	:      Optional[Any]     							=					self.retrieval_workers[random.randint(0     ,	len(self.retrieval_workers      ) - 1      )]
              UpperCAmelCase__							,							UpperCAmelCase__	:      List[Any]     							=					ray.get(random_worker.retrieve.remote(UpperCamelCase_     ,	UpperCamelCase_      )      )
          else:
              UpperCAmelCase__							,							UpperCAmelCase__	:      Dict     							=					self._main_retrieve(UpperCamelCase_     ,	UpperCamelCase_      )
          return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(UpperCamelCase_      )



      @classmethod
      def  __snake_case							(   cls     ,	UpperCamelCase_     ,	UpperCamelCase_=None     ,	**UpperCamelCase_      ):
          return super(UpperCamelCase_     ,	cls      ).get_tokenizers(UpperCamelCase_     ,	UpperCamelCase_     ,	**UpperCamelCase_      )



      @classmethod
      def  __snake_case							(   cls     ,	UpperCamelCase_     ,	UpperCamelCase_     ,	UpperCamelCase_=None     ,	**UpperCamelCase_      ):
          UpperCAmelCase__	:      Union[str, Any]     							=					kwargs.pop('config'     ,	UpperCamelCase_      ) or RagConfig.from_pretrained(UpperCamelCase_     ,	**UpperCamelCase_      )
          UpperCAmelCase__	:      Union[str, Any]     							=					RagTokenizer.from_pretrained(UpperCamelCase_     ,	config=UpperCamelCase_      )
          UpperCAmelCase__	:      List[Any]     							=					rag_tokenizer.question_encoder
          UpperCAmelCase__	:      Optional[Any]     							=					rag_tokenizer.generator
          if indexed_dataset is not None:
              UpperCAmelCase__	:      List[str]     							=					'custom'
              UpperCAmelCase__	:      Tuple     							=					CustomHFIndex(config.retrieval_vector_size     ,	UpperCamelCase_      )
          else:
              UpperCAmelCase__	:      Optional[Any]     							=					cls._build_index(UpperCamelCase_      )
          return cls(
              UpperCamelCase_     ,	question_encoder_tokenizer=UpperCamelCase_     ,	generator_tokenizer=UpperCamelCase_     ,	retrieval_workers=UpperCamelCase_     ,	index=UpperCamelCase_     ,	)

