


import inspect
import unittest

from transformers import ResNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
		import torch
		from torch import nn

		from transformers import ResNetBackbone, ResNetForImageClassification, ResNetModel
		from transformers.models.resnet.modeling_resnet import RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor





class    _a				:




	def __init__(   self		:						List[Any]						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						Any=3						,   _SCREAMING_SNAKE_CASE		:						Any=32						,   _SCREAMING_SNAKE_CASE		:						str=3						,   _SCREAMING_SNAKE_CASE		:						int=10						,   _SCREAMING_SNAKE_CASE		:						Optional[int]=[10, 20, 30, 40]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]=[1, 1, 2, 1]						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]=True						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]=True						,   _SCREAMING_SNAKE_CASE		:						List[Any]="relu"						,   _SCREAMING_SNAKE_CASE		:						List[str]=3						,   _SCREAMING_SNAKE_CASE		:						int=None						,   _SCREAMING_SNAKE_CASE		:						List[str]=["stage2", "stage3", "stage4"]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]=[2, 3, 4]						,   )->			Dict:
			lowerCAmelCase__  :     Tuple       	=   parent
			lowerCAmelCase__  :     List[str]       	=   batch_size
			lowerCAmelCase__  :     Union[str, Any]       	=   image_size
			lowerCAmelCase__  :     Any       	=   num_channels
			lowerCAmelCase__  :     Optional[int]       	=   embeddings_size
			lowerCAmelCase__  :     int       	=   hidden_sizes
			lowerCAmelCase__  :     Dict       	=   depths
			lowerCAmelCase__  :     int       	=   is_training
			lowerCAmelCase__  :     Dict       	=   use_labels
			lowerCAmelCase__  :     Any       	=   hidden_act
			lowerCAmelCase__  :     int       	=   num_labels
			lowerCAmelCase__  :     List[str]       	=   scope
			lowerCAmelCase__  :     int       	=   len(_lowerCAmelCase						)
			lowerCAmelCase__  :     str       	=   out_features
			lowerCAmelCase__  :     Any       	=   out_indices




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						)->			Union[str, Any]:
			lowerCAmelCase__  :     Tuple       	=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

			lowerCAmelCase__  :     Optional[Any]       	=   None
			if self.use_labels:
					lowerCAmelCase__  :     int       	=   ids_tensor([self.batch_size]						,   self.num_labels						)

			lowerCAmelCase__  :     List[Any]       	=   self.get_config()

			return config, pixel_values, labels




	def 				UpperCAmelCase__(   self		:						Tuple						)->			Any:
			return ResNetConfig(
			    num_channels=self.num_channels						,   embeddings_size=self.embeddings_size						,   hidden_sizes=self.hidden_sizes						,   depths=self.depths						,   hidden_act=self.hidden_act						,   num_labels=self.num_labels						,   out_features=self.out_features						,   out_indices=self.out_indices						,   )




	def 				UpperCAmelCase__(   self		:						List[Any]						,   _SCREAMING_SNAKE_CASE		:						List[Any]						,   _SCREAMING_SNAKE_CASE		:						List[str]						,   _SCREAMING_SNAKE_CASE		:						Dict						)->			Optional[int]:
			lowerCAmelCase__  :     Union[str, Any]       	=   ResNetModel(config=_lowerCAmelCase						)
			model.to(_lowerCAmelCase						)
			model.eval()
			lowerCAmelCase__  :     List[str]       	=   model(_lowerCAmelCase						)
			# expected last hidden states: B, C, H // 32, W // 32
			self.parent.assertEqual(
			    result.last_hidden_state.shape						,   (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)						,   )




	def 				UpperCAmelCase__(   self		:						str						,   _SCREAMING_SNAKE_CASE		:						List[str]						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						,   _SCREAMING_SNAKE_CASE		:						int						)->			Tuple:
			lowerCAmelCase__  :     Tuple       	=   self.num_labels
			lowerCAmelCase__  :     Union[str, Any]       	=   ResNetForImageClassification(_lowerCAmelCase						)
			model.to(_lowerCAmelCase						)
			model.eval()
			lowerCAmelCase__  :     List[str]       	=   model(_lowerCAmelCase						,   labels=_lowerCAmelCase						)
			self.parent.assertEqual(result.logits.shape						,   (self.batch_size, self.num_labels)						)




	def 				UpperCAmelCase__(   self		:						Optional[int]						,   _SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						int						,   _SCREAMING_SNAKE_CASE		:						Union[str, Any]						)->			int:
			lowerCAmelCase__  :     str       	=   ResNetBackbone(config=_lowerCAmelCase						)
			model.to(_lowerCAmelCase						)
			model.eval()
			lowerCAmelCase__  :     int       	=   model(_lowerCAmelCase						)

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps						)						,   len(config.out_features						)						)
			self.parent.assertListEqual(list(result.feature_maps[0].shape						)						,   [self.batch_size, self.hidden_sizes[1], 4, 4]						)

			# verify channels
			self.parent.assertEqual(len(model.channels						)						,   len(config.out_features						)						)
			self.parent.assertListEqual(model.channels						,   config.hidden_sizes[1:]						)

			# verify backbone works with out_features=None
			lowerCAmelCase__  :     Optional[Any]       	=   None
			lowerCAmelCase__  :     List[str]       	=   ResNetBackbone(config=_lowerCAmelCase						)
			model.to(_lowerCAmelCase						)
			model.eval()
			lowerCAmelCase__  :     Optional[Any]       	=   model(_lowerCAmelCase						)

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps						)						,   1						)
			self.parent.assertListEqual(list(result.feature_maps[0].shape						)						,   [self.batch_size, self.hidden_sizes[-1], 1, 1]						)

			# verify channels
			self.parent.assertEqual(len(model.channels						)						,   1						)
			self.parent.assertListEqual(model.channels						,   [config.hidden_sizes[-1]]						)




	def 				UpperCAmelCase__(   self		:						int						)->			Optional[int]:
			lowerCAmelCase__  :     Any       	=   self.prepare_config_and_inputs()
			lowerCAmelCase__	,      lowerCAmelCase__	,      lowerCAmelCase__  :     List[Any]       	=   config_and_inputs
			lowerCAmelCase__  :     Dict       	=   {'''pixel_values''': pixel_values}
			return config, inputs_dict





@require_torch
class    _a				(							_lowercase	,		_lowercase	,		unittest.TestCase):
	_a    :    int        =				(
	    (
	        ResNetModel,
	        ResNetForImageClassification,
	        ResNetBackbone,
	    )
	    if is_torch_available()
	    else ()
	)
	_a    :    Tuple        =				(
	    {'''feature-extraction''': ResNetModel, '''image-classification''': ResNetForImageClassification}
	    if is_torch_available()
	    else {}
	)

	_a    :    List[Any]        =				True
	_a    :    Any        =				False
	_a    :    Tuple        =				False
	_a    :    List[Any]        =				False
	_a    :    Union[str, Any]        =				False




	def 				UpperCAmelCase__(   self		:						str						)->			Dict:
			lowerCAmelCase__  :     Union[str, Any]       	=   ResNetModelTester(self						)
			lowerCAmelCase__  :     Any       	=   ConfigTester(self						,   config_class=_lowerCAmelCase						,   has_text_modality=_lowerCAmelCase						)




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						)->			List[str]:
			self.create_and_test_config_common_properties()
			self.config_tester.create_and_test_config_to_json_string()
			self.config_tester.create_and_test_config_to_json_file()
			self.config_tester.create_and_test_config_from_and_save_pretrained()
			self.config_tester.create_and_test_config_with_num_labels()
			self.config_tester.check_config_can_be_init_without_params()
			self.config_tester.check_config_arguments_init()




	def 				UpperCAmelCase__(   self		:						Union[str, Any]						)->			Optional[int]:
			return




	@unittest.skip(reason='''ResNet does not use inputs_embeds'''						)
	def 				UpperCAmelCase__(   self		:						List[str]						)->			Tuple:
			pass




	@unittest.skip(reason='''ResNet does not support input and output embeddings'''						)
	def 				UpperCAmelCase__(   self		:						Optional[int]						)->			Any:
			pass




	def 				UpperCAmelCase__(   self		:						Optional[Any]						)->			Tuple:
			lowerCAmelCase__	,      lowerCAmelCase__  :     int       	=   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					lowerCAmelCase__  :     Dict       	=   model_class(_lowerCAmelCase						)
					lowerCAmelCase__  :     Optional[Any]       	=   inspect.signature(model.forward						)
					# signature.parameters is an OrderedDict => so arg_names order is deterministic
					lowerCAmelCase__  :     int       	=   [*signature.parameters.keys()]

					lowerCAmelCase__  :     Union[str, Any]       	=   ['''pixel_values''']
					self.assertListEqual(arg_names[:1]						,   _lowerCAmelCase						)




	def 				UpperCAmelCase__(   self		:						str						)->			Any:
			lowerCAmelCase__  :     Dict       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*_lowerCAmelCase						)




	def 				UpperCAmelCase__(   self		:						Optional[int]						)->			int:
			lowerCAmelCase__  :     List[str]       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_backbone(*_lowerCAmelCase						)




	def 				UpperCAmelCase__(   self		:						str						)->			Dict:
			lowerCAmelCase__	,      lowerCAmelCase__  :     List[str]       	=   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					lowerCAmelCase__  :     Optional[int]       	=   model_class(config=_lowerCAmelCase						)
					for name, module in model.named_modules():
							if isinstance(_lowerCAmelCase						,   (nn.BatchNormad, nn.GroupNorm)						):
									self.assertTrue(
									    torch.all(module.weight == 1						)						,   msg=F'Parameter {name} of model {model_class} seems not properly initialized'						,   )
									self.assertTrue(
									    torch.all(module.bias == 0						)						,   msg=F'Parameter {name} of model {model_class} seems not properly initialized'						,   )




	def 				UpperCAmelCase__(   self		:						Tuple						)->			List[str]:
			def check_hidden_states_output(_SCREAMING_SNAKE_CASE		:						Optional[Any]						,   _SCREAMING_SNAKE_CASE		:						Any						,   _SCREAMING_SNAKE_CASE		:						List[Any]						):
					lowerCAmelCase__  :     List[str]       	=   model_class(_lowerCAmelCase						)
					model.to(_lowerCAmelCase						)
					model.eval()

					with torch.no_grad():
							lowerCAmelCase__  :     List[str]       	=   model(**self._prepare_for_class(_lowerCAmelCase						,   _lowerCAmelCase						)						)

					lowerCAmelCase__  :     int       	=   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

					lowerCAmelCase__  :     Tuple       	=   self.model_tester.num_stages
					self.assertEqual(len(_lowerCAmelCase						)						,   expected_num_stages + 1						)

					# ResNet's feature maps are of shape (batch_size, num_channels, height, width)
					self.assertListEqual(
					    list(hidden_states[0].shape[-2:]						)						,   [self.model_tester.image_size // 4, self.model_tester.image_size // 4]						,   )

			lowerCAmelCase__	,      lowerCAmelCase__  :     Optional[int]       	=   self.model_tester.prepare_config_and_inputs_for_common()
			lowerCAmelCase__  :     List[str]       	=   ['''basic''', '''bottleneck''']
			for model_class in self.all_model_classes:
					for layer_type in layers_type:
							lowerCAmelCase__  :     Optional[int]       	=   layer_type
							lowerCAmelCase__  :     Tuple       	=   True
							check_hidden_states_output(_lowerCAmelCase						,   _lowerCAmelCase						,   _lowerCAmelCase						)

							# check that output_hidden_states also work using config
							del inputs_dict["output_hidden_states"]
							lowerCAmelCase__  :     Tuple       	=   True

							check_hidden_states_output(_lowerCAmelCase						,   _lowerCAmelCase						,   _lowerCAmelCase						)




	@unittest.skip(reason='''ResNet does not use feedforward chunking'''						)
	def 				UpperCAmelCase__(   self		:						Optional[Any]						)->			Optional[Any]:
			pass




	def 				UpperCAmelCase__(   self		:						Optional[Any]						)->			Optional[Any]:
			lowerCAmelCase__  :     Union[str, Any]       	=   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*_lowerCAmelCase						)




	@slow
	def 				UpperCAmelCase__(   self		:						Any						)->			Any:
			for model_name in RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCAmelCase__  :     Dict       	=   ResNetModel.from_pretrained(_lowerCAmelCase						)
					self.assertIsNotNone(_lowerCAmelCase						)






def 	lowerCamelCase_       (			):




		"""simple docstring"""






		lowerCAmelCase__  :     Any       	=   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''			)
		return image





@require_torch
@require_vision
class    _a				(							unittest.TestCase):




	@cached_property
	def 				UpperCAmelCase__(   self		:						List[str]						)->			Any:
			return (
			    AutoImageProcessor.from_pretrained(RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)
			    if is_vision_available()
			    else None
			)




	@slow
	def 				UpperCAmelCase__(   self		:						Dict						)->			Dict:
			lowerCAmelCase__  :     Tuple       	=   ResNetForImageClassification.from_pretrained(RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]						).to(_lowerCAmelCase						)

			lowerCAmelCase__  :     int       	=   self.default_image_processor
			lowerCAmelCase__  :     Union[str, Any]       	=   prepare_img()
			lowerCAmelCase__  :     str       	=   image_processor(images=_lowerCAmelCase						,   return_tensors='''pt'''						).to(_lowerCAmelCase						)

			# forward pass
			with torch.no_grad():
					lowerCAmelCase__  :     Union[str, Any]       	=   model(**_lowerCAmelCase						)

			# verify the logits
			lowerCAmelCase__  :     List[Any]       	=   torch.Size((1, 1000)						)
			self.assertEqual(outputs.logits.shape						,   _lowerCAmelCase						)

			lowerCAmelCase__  :     List[str]       	=   torch.tensor([-11.1069, -9.7877, -8.3777]						).to(_lowerCAmelCase						)

			self.assertTrue(torch.allclose(outputs.logits[0, :3]						,   _lowerCAmelCase						,   atol=1E-4						)						)





@require_torch
class    _a				(							_lowercase	,		unittest.TestCase):
	_a    :    Dict        =				(ResNetBackbone,) if is_torch_available() else ()
	_a    :    Optional[int]        =				False
	_a    :    str        =				ResNetConfig




	def 				UpperCAmelCase__(   self		:						List[Any]						)->			int:
			lowerCAmelCase__  :     Optional[Any]       	=   ResNetModelTester(self						)






import json
import os
import shutil
import sys
import tempfile
import unittest
from unittest import TestCase
from unittest.mock import patch

import faiss
import numpy as np
from datasets import Dataset

from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
from transformers.integrations import is_ray_available
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_ray


sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

if is_torch_available():
       from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
else:
       _UpperCamelCase:					List[Any]     =None

if is_ray_available():
       import ray  # noqa: E402 # isort:skip
       from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
       _UpperCamelCase:					List[str]     =None
       _UpperCamelCase:					Tuple     =None
       _UpperCamelCase:					Any     =None



def 	_a  (					__SCREAMING_SNAKE_CASE					:  Optional[int]    ):






  """simple docstring"""




  if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
    _lowerCAmelCase						=		unittest.skip('test requires Datasets, Faiss, psutil'    )(__SCREAMING_SNAKE_CASE    )
  return test_case





@require_distributed_retrieval
class 						__lowercase(     SCREAMING_SNAKE_CASE   ):


       """simple docstring"""







       def SCREAMING_SNAKE_CASE_				(     self					:			List[str]       )					->   Optional[Any]:
         _lowerCAmelCase						=		tempfile.mkdtemp()
         _lowerCAmelCase						=		8

         # DPR tok
         _lowerCAmelCase						=		[
             '[UNK]',
             '[CLS]',
             '[SEP]',
             '[PAD]',
             '[MASK]',
             'want',
             '##want',
             '##ed',
             'wa',
             'un',
             'runn',
             '##ing',
             ',',
             'low',
             'lowest',
         ]
         _lowerCAmelCase						=		os.path.join(self.tmpdirname  ,      'dpr_tokenizer'       )
         os.makedirs(_lowerCAmelCase  ,      exist_ok=_lowerCAmelCase       )
         _lowerCAmelCase						=		os.path.join(_lowerCAmelCase  ,      DPR_VOCAB_FILES_NAMES['vocab_file']       )
         with open(self.vocab_file  ,      'w'  ,      encoding='utf-8'       ) as vocab_writer:
           vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]       )       )

         # BART tok
         _lowerCAmelCase						=		[
             'l',
             'o',
             'w',
             'e',
             'r',
             's',
             't',
             'i',
             'd',
             'n',
             '\u0120',
             '\u0120l',
             '\u0120n',
             '\u0120lo',
             '\u0120low',
             'er',
             '\u0120lowest',
             '\u0120newer',
             '\u0120wider',
             '<unk>',
         ]
         _lowerCAmelCase						=		dict(zip(_lowerCAmelCase  ,      range(len(_lowerCAmelCase       )       )       )       )
         _lowerCAmelCase						=		['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
         _lowerCAmelCase						=		{'unk_token': '<unk>'}

         _lowerCAmelCase						=		os.path.join(self.tmpdirname  ,      'bart_tokenizer'       )
         os.makedirs(_lowerCAmelCase  ,      exist_ok=_lowerCAmelCase       )
         _lowerCAmelCase						=		os.path.join(_lowerCAmelCase  ,      BART_VOCAB_FILES_NAMES['vocab_file']       )
         _lowerCAmelCase						=		os.path.join(_lowerCAmelCase  ,      BART_VOCAB_FILES_NAMES['merges_file']       )
         with open(self.vocab_file  ,      'w'  ,      encoding='utf-8'       ) as fp:
           fp.write(json.dumps(_lowerCAmelCase       ) + '\n'       )
         with open(self.merges_file  ,      'w'  ,      encoding='utf-8'       ) as fp:
           fp.write('\n'.join(_lowerCAmelCase       )       )



       def SCREAMING_SNAKE_CASE_				(     self					:			Tuple       )					->   DPRQuestionEncoderTokenizer:
         return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,      'dpr_tokenizer'       )       )



       def SCREAMING_SNAKE_CASE_				(     self					:			Any       )					->   BartTokenizer:
         return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,      'bart_tokenizer'       )       )



       def SCREAMING_SNAKE_CASE_				(     self					:			Dict       )					->   int:
         shutil.rmtree(self.tmpdirname       )



       def SCREAMING_SNAKE_CASE_				(     self					:			str       )					->   str:
         _lowerCAmelCase						=		Dataset.from_dict(
             {
                 'id': ['0', '1'],
                 'text': ['foo', 'bar'],
                 'title': ['Foo', 'Bar'],
                 'embeddings': [np.ones(self.retrieval_vector_size       ), 2 * np.ones(self.retrieval_vector_size       )],
             }       )
         dataset.add_faiss_index('embeddings'  ,      string_factory='Flat'  ,      metric_type=faiss.METRIC_INNER_PRODUCT       )
         return dataset



       def SCREAMING_SNAKE_CASE_				(     self					:			int  ,      _lowerCAmelCase					:			bool  ,      _lowerCAmelCase					:			Tuple=1_2345       )					->   RagPyTorchDistributedRetriever:
         _lowerCAmelCase						=		self.get_dummy_dataset()
         _lowerCAmelCase						=		RagConfig(
             retrieval_vector_size=self.retrieval_vector_size  ,      question_encoder=DPRConfig().to_dict()  ,      generator=BartConfig().to_dict()  ,      )
         with patch('transformers.models.rag.retrieval_rag.load_dataset'       ) as mock_load_dataset:
           _lowerCAmelCase						=		dataset
           _lowerCAmelCase						=		RagPyTorchDistributedRetriever(
               _lowerCAmelCase  ,      question_encoder_tokenizer=self.get_dpr_tokenizer()  ,      generator_tokenizer=self.get_bart_tokenizer()  ,      )
           if init_retrieval:
             retriever.init_retrieval(_lowerCAmelCase       )
         return retriever



       def SCREAMING_SNAKE_CASE_				(     self					:			Any  ,      _lowerCAmelCase					:			bool       )					->   RagRayDistributedRetriever:
         # Have to run in local mode because sys.path modifications at top of
         # file are not propogated to remote workers.
         # https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
         ray.init(local_mode=_lowerCAmelCase       )
         _lowerCAmelCase						=		RagConfig(
             retrieval_vector_size=self.retrieval_vector_size  ,      question_encoder=DPRConfig().to_dict()  ,      generator=BartConfig().to_dict()  ,      )
         _lowerCAmelCase						=		ray.remote(_lowerCAmelCase       )
         _lowerCAmelCase						=		[remote_cls.remote() for _ in range(1       )]
         with patch('transformers.models.rag.retrieval_rag.load_dataset'       ) as mock_load_dataset:
           _lowerCAmelCase						=		self.get_dummy_dataset()
           _lowerCAmelCase						=		RagRayDistributedRetriever(
               _lowerCAmelCase  ,      question_encoder_tokenizer=self.get_dpr_tokenizer()  ,      generator_tokenizer=self.get_bart_tokenizer()  ,      retrieval_workers=_lowerCAmelCase  ,      )
           if init_retrieval:
             retriever.init_retrieval()
         return retriever



       def SCREAMING_SNAKE_CASE_				(     self					:			Optional[Any]  ,      _lowerCAmelCase					:			bool  ,      _lowerCAmelCase					:			bool  ,      _lowerCAmelCase					:			List[Any]=1_2345       )					->   List[Any]:
         _lowerCAmelCase						=		self.get_dummy_dataset()
         _lowerCAmelCase						=		RagConfig(
             retrieval_vector_size=self.retrieval_vector_size  ,      question_encoder=DPRConfig().to_dict()  ,      generator=BartConfig().to_dict()  ,      index_name='custom'  ,      )
         if from_disk:
           _lowerCAmelCase						=		os.path.join(self.tmpdirname  ,      'dataset'       )
           _lowerCAmelCase						=		os.path.join(self.tmpdirname  ,      'index.faiss'       )
           dataset.get_index('embeddings'       ).save(os.path.join(self.tmpdirname  ,      'index.faiss'       )       )
           dataset.drop_index('embeddings'       )
           dataset.save_to_disk(os.path.join(self.tmpdirname  ,      'dataset'       )       )
           del dataset
           _lowerCAmelCase						=		RagPyTorchDistributedRetriever(
               _lowerCAmelCase  ,      question_encoder_tokenizer=self.get_dpr_tokenizer()  ,      generator_tokenizer=self.get_bart_tokenizer()  ,      )
         else:
           _lowerCAmelCase						=		RagPyTorchDistributedRetriever(
               _lowerCAmelCase  ,      question_encoder_tokenizer=self.get_dpr_tokenizer()  ,      generator_tokenizer=self.get_bart_tokenizer()  ,      index=CustomHFIndex(config.retrieval_vector_size  ,      _lowerCAmelCase       )  ,      )
         if init_retrieval:
           retriever.init_retrieval(_lowerCAmelCase       )
         return retriever



       def SCREAMING_SNAKE_CASE_				(     self					:			Union[str, Any]  ,      _lowerCAmelCase					:			bool  ,      _lowerCAmelCase					:			bool       )					->   int:
         # Have to run in local mode because sys.path modifications at top of
         # file are not propogated to remote workers.
         # https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
         ray.init(local_mode=_lowerCAmelCase       )
         _lowerCAmelCase						=		self.get_dummy_dataset()
         _lowerCAmelCase						=		RagConfig(
             retrieval_vector_size=self.retrieval_vector_size  ,      question_encoder=DPRConfig().to_dict()  ,      generator=BartConfig().to_dict()  ,      index_name='custom'  ,      )
         _lowerCAmelCase						=		ray.remote(_lowerCAmelCase       )
         _lowerCAmelCase						=		[remote_cls.remote() for _ in range(1       )]
         if from_disk:
           _lowerCAmelCase						=		os.path.join(self.tmpdirname  ,      'dataset'       )
           _lowerCAmelCase						=		os.path.join(self.tmpdirname  ,      'index.faiss'       )
           dataset.get_index('embeddings'       ).save(os.path.join(self.tmpdirname  ,      'index.faiss'       )       )
           dataset.drop_index('embeddings'       )
           dataset.save_to_disk(os.path.join(self.tmpdirname  ,      'dataset'       )       )
           del dataset
           _lowerCAmelCase						=		RagRayDistributedRetriever(
               _lowerCAmelCase  ,      question_encoder_tokenizer=self.get_dpr_tokenizer()  ,      generator_tokenizer=self.get_bart_tokenizer()  ,      retrieval_workers=_lowerCAmelCase  ,      index=CustomHFIndex.load_from_disk(
                   vector_size=config.retrieval_vector_size  ,      dataset_path=config.passages_path  ,      index_path=config.index_path  ,      )  ,      )
         else:
           _lowerCAmelCase						=		RagRayDistributedRetriever(
               _lowerCAmelCase  ,      question_encoder_tokenizer=self.get_dpr_tokenizer()  ,      generator_tokenizer=self.get_bart_tokenizer()  ,      retrieval_workers=_lowerCAmelCase  ,      index=CustomHFIndex(config.retrieval_vector_size  ,      _lowerCAmelCase       )  ,      )
         if init_retrieval:
           retriever.init_retrieval()
         return retriever



       def SCREAMING_SNAKE_CASE_				(     self					:			str  ,      _lowerCAmelCase					:			RagRetriever  ,      _lowerCAmelCase					:			np.array  ,      _lowerCAmelCase					:			int       )					->   None:
         _lowerCAmelCase	,       _lowerCAmelCase	,       _lowerCAmelCase						=		retriever.retrieve(_lowerCAmelCase  ,      n_docs=_lowerCAmelCase       )
         self.assertEqual(retrieved_doc_embeds.shape  ,      (2, n_docs, self.retrieval_vector_size)       )
         self.assertEqual(len(_lowerCAmelCase       )  ,      2       )
         self.assertEqual(sorted(doc_dicts[0]       )  ,      ['embeddings', 'id', 'text', 'title']       )
         self.assertEqual(len(doc_dicts[0]['id']       )  ,      _lowerCAmelCase       )
         self.assertEqual(doc_dicts[0]['id'][0]  ,      '1'       )  # max inner product is reached with second doc
         self.assertEqual(doc_dicts[1]['id'][0]  ,      '0'       )  # max inner product is reached with first doc
         self.assertListEqual(doc_ids.tolist()  ,      [[1], [0]]       )



       def SCREAMING_SNAKE_CASE_				(     self					:			Any       )					->   Optional[int]:
         _lowerCAmelCase						=		1
         _lowerCAmelCase						=		np.array(
             [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]  ,      dtype=np.floataa       )

         self.distributed_retriever_check(
             self.get_dummy_pytorch_distributed_retriever(init_retrieval=_lowerCAmelCase       )  ,      _lowerCAmelCase  ,      _lowerCAmelCase       )



       def SCREAMING_SNAKE_CASE_				(     self					:			Tuple       )					->   Any:
         _lowerCAmelCase						=		1
         _lowerCAmelCase						=		np.array(
             [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]  ,      dtype=np.floataa       )

         self.distributed_retriever_check(
             self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=_lowerCAmelCase  ,      from_disk=_lowerCAmelCase       )  ,      _lowerCAmelCase  ,      _lowerCAmelCase  ,      )



       def SCREAMING_SNAKE_CASE_				(     self					:			int       )					->   Optional[Any]:
         _lowerCAmelCase						=		1
         _lowerCAmelCase						=		np.array(
             [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]  ,      dtype=np.floataa       )

         self.distributed_retriever_check(
             self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=_lowerCAmelCase  ,      from_disk=_lowerCAmelCase       )  ,      _lowerCAmelCase  ,      _lowerCAmelCase  ,      )



       @require_ray
       def SCREAMING_SNAKE_CASE_				(     self					:			Dict       )					->   Any:
         _lowerCAmelCase						=		1
         _lowerCAmelCase						=		np.array(
             [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]  ,      dtype=np.floataa       )

         self.distributed_retriever_check(
             self.get_dummy_ray_distributed_retriever(init_retrieval=_lowerCAmelCase       )  ,      _lowerCAmelCase  ,      _lowerCAmelCase       )
         ray.shutdown()



       @require_ray
       def SCREAMING_SNAKE_CASE_				(     self					:			List[str]       )					->   Optional[Any]:
         _lowerCAmelCase						=		1
         _lowerCAmelCase						=		np.array(
             [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]  ,      dtype=np.floataa       )
         with self.assertRaises(_lowerCAmelCase       ):
           self.distributed_retriever_check(
               self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=_lowerCAmelCase  ,      from_disk=_lowerCAmelCase       )  ,      _lowerCAmelCase  ,      _lowerCAmelCase  ,      )
         ray.shutdown()






       @require_ray
       def SCREAMING_SNAKE_CASE_				(     self					:			str       )					->   List[Any]:
         _lowerCAmelCase						=		1
         _lowerCAmelCase						=		np.array(
             [np.ones(self.retrieval_vector_size       ), -np.ones(self.retrieval_vector_size       )]  ,      dtype=np.floataa       )

         self.distributed_retriever_check(
             self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=_lowerCAmelCase  ,      from_disk=_lowerCAmelCase       )  ,      _lowerCAmelCase  ,      _lowerCAmelCase       )
         ray.shutdown()
