





"""simple docstring"""





import inspect
import unittest

from transformers import DataaVecVisionConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
    import torch
    from torch import nn

    from transformers import (
        MODEL_MAPPING,
        DataaVecVisionForImageClassification,
        DataaVecVisionForSemanticSegmentation,
        DataaVecVisionModel,
    )
    from transformers.models.dataavec.modeling_dataavec_vision import DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import BeitImageProcessor





class 			snake_case		:

      def __init__(			self   :List[Any]						,							_lowerCamelCase   :List[Any]						,							_lowerCamelCase   :Dict=1_0_0						,							_lowerCamelCase   :Any=1_3						,							_lowerCamelCase   :Optional[int]=3_0						,							_lowerCamelCase   :Union[str, Any]=2						,							_lowerCamelCase   :Tuple=3						,							_lowerCamelCase   :int=True						,							_lowerCamelCase   :Optional[Any]=True						,							_lowerCamelCase   :int=3_2						,							_lowerCamelCase   :List[Any]=4						,							_lowerCamelCase   :Dict=4						,							_lowerCamelCase   :List[Any]=3_7						,							_lowerCamelCase   :Optional[Any]="gelu"						,							_lowerCamelCase   :List[str]=0.1						,							_lowerCamelCase   :Tuple=0.1						,							_lowerCamelCase   :List[str]=1_0						,							_lowerCamelCase   :List[Any]=0.0_2						,							_lowerCamelCase   :Optional[int]=3						,							_lowerCamelCase   :Tuple=None						,							_lowerCamelCase   :Union[str, Any]=[0, 1, 2, 3]						,							):
       __SCREAMING_SNAKE_CASE :			Any												=       parent
       __SCREAMING_SNAKE_CASE :			Dict												=       1_0_0
       __SCREAMING_SNAKE_CASE :			str												=       batch_size
       __SCREAMING_SNAKE_CASE :			List[Any]												=       image_size
       __SCREAMING_SNAKE_CASE :			Dict												=       patch_size
       __SCREAMING_SNAKE_CASE :			Tuple												=       num_channels
       __SCREAMING_SNAKE_CASE :			Dict												=       is_training
       __SCREAMING_SNAKE_CASE :			Tuple												=       use_labels
       __SCREAMING_SNAKE_CASE :			List[str]												=       hidden_size
       __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       num_hidden_layers
       __SCREAMING_SNAKE_CASE :			Dict												=       num_attention_heads
       __SCREAMING_SNAKE_CASE :			Dict												=       intermediate_size
       __SCREAMING_SNAKE_CASE :			Any												=       hidden_act
       __SCREAMING_SNAKE_CASE :			Optional[Any]												=       hidden_dropout_prob
       __SCREAMING_SNAKE_CASE :			List[Any]												=       attention_probs_dropout_prob
       __SCREAMING_SNAKE_CASE :			List[Any]												=       type_sequence_label_size
       __SCREAMING_SNAKE_CASE :			Any												=       initializer_range
       __SCREAMING_SNAKE_CASE :			Dict												=       scope
       __SCREAMING_SNAKE_CASE :			Tuple												=       out_indices
       __SCREAMING_SNAKE_CASE :			int												=       num_labels

       # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
       __SCREAMING_SNAKE_CASE :			Any												=       (image_size // patch_size) ** 2
       __SCREAMING_SNAKE_CASE :			List[str]												=       num_patches + 1

      def 						SCREAMING_SNAKE_CASE_			(			self   :str						):
       __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

       __SCREAMING_SNAKE_CASE :			Any												=       None
       __SCREAMING_SNAKE_CASE :			Tuple												=       None
       if self.use_labels:
        __SCREAMING_SNAKE_CASE :			List[Any]												=       ids_tensor([self.batch_size]						,							self.type_sequence_label_size						)
        __SCREAMING_SNAKE_CASE :			Optional[int]												=       ids_tensor([self.batch_size, self.image_size, self.image_size]						,							self.num_labels						)

       __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       self.get_config()

       return config, pixel_values, labels, pixel_labels

      def 						SCREAMING_SNAKE_CASE_			(			self   :Dict						):
       return DataaVecVisionConfig(
           vocab_size=self.vocab_size						,							image_size=self.image_size						,							patch_size=self.patch_size						,							num_channels=self.num_channels						,							hidden_size=self.hidden_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						,							is_decoder=_lowerCamelCase						,							initializer_range=self.initializer_range						,							out_indices=self.out_indices						,							)

      def 						SCREAMING_SNAKE_CASE_			(			self   :Union[str, Any]						,							_lowerCamelCase   :Optional[Any]						,							_lowerCamelCase   :Tuple						,							_lowerCamelCase   :Dict						,							_lowerCamelCase   :List[Any]						):
       __SCREAMING_SNAKE_CASE :			int												=       DataaVecVisionModel(config=_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       model(_lowerCamelCase						)
       # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
       __SCREAMING_SNAKE_CASE :			Tuple												=       (self.image_size // self.patch_size) ** 2
       self.parent.assertEqual(result.last_hidden_state.shape						,							(self.batch_size, num_patches + 1, self.hidden_size)						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :Tuple						,							_lowerCamelCase   :Any						,							_lowerCamelCase   :List[Any]						,							_lowerCamelCase   :int						,							_lowerCamelCase   :Optional[Any]						):
       __SCREAMING_SNAKE_CASE :			Any												=       self.type_sequence_label_size
       __SCREAMING_SNAKE_CASE :			List[Any]												=       DataaVecVisionForImageClassification(_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       __SCREAMING_SNAKE_CASE :			List[str]												=       model(_lowerCamelCase						,							labels=_lowerCamelCase						)
       self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.type_sequence_label_size)						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :Any						,							_lowerCamelCase   :List[str]						,							_lowerCamelCase   :int						,							_lowerCamelCase   :Union[str, Any]						,							_lowerCamelCase   :Dict						):
       __SCREAMING_SNAKE_CASE :			Tuple												=       self.num_labels
       __SCREAMING_SNAKE_CASE :			Tuple												=       DataaVecVisionForSemanticSegmentation(_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       __SCREAMING_SNAKE_CASE :			Any												=       model(_lowerCamelCase						)
       self.parent.assertEqual(
           result.logits.shape						,							(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)						)
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       model(_lowerCamelCase						,							labels=_lowerCamelCase						)
       self.parent.assertEqual(
           result.logits.shape						,							(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)						)



      def 						SCREAMING_SNAKE_CASE_			(			self   :List[str]						):
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       self.prepare_config_and_inputs()
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			List[str]												=       config_and_inputs
       __SCREAMING_SNAKE_CASE :			Any												=       {'''pixel_values''': pixel_values}
       return config, inputs_dict







@require_torch
class 			snake_case		(   __UpperCAmelCase ,     __UpperCAmelCase ,     unittest.TestCase   ):
      lowerCamelCase__						= (
          (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
          if is_torch_available()
          else ()
      )
      lowerCamelCase__						= (
          {
              '''feature-extraction''': DataaVecVisionModel,
              '''image-classification''': DataaVecVisionForImageClassification,
              '''image-segmentation''': DataaVecVisionForSemanticSegmentation,
          }
          if is_torch_available()
          else {}
      )

      lowerCamelCase__						= False
      lowerCamelCase__						= False
      lowerCamelCase__						= False

      def 						SCREAMING_SNAKE_CASE_			(			self   :List[Any]						):
       __SCREAMING_SNAKE_CASE :			Any												=       DataaVecVisionModelTester(self						)
       __SCREAMING_SNAKE_CASE :			int												=       ConfigTester(
           self						,							config_class=_lowerCamelCase						,							has_text_modality=_lowerCamelCase						,							hidden_size=3_7						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :int						):
       self.config_tester.run_common_tests()

      def 						SCREAMING_SNAKE_CASE_			(			self   :List[Any]						):
       # Data2VecVision does not use inputs_embeds
       pass

      @require_torch_multi_gpu
      @unittest.skip(
          reason='''Data2VecVision has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''						)
      def 						SCREAMING_SNAKE_CASE_			(			self   :int						):
       pass

      def 						SCREAMING_SNAKE_CASE_			(			self   :str						):
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			Tuple												=       self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        __SCREAMING_SNAKE_CASE :			Any												=       model_class(_lowerCamelCase						)
        self.assertIsInstance(model.get_input_embeddings()						,							(nn.Module)						)
        __SCREAMING_SNAKE_CASE :			Tuple												=       model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(_lowerCamelCase						,							nn.Linear						)						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :Dict						):
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			Dict												=       self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        __SCREAMING_SNAKE_CASE :			str												=       model_class(_lowerCamelCase						)
        __SCREAMING_SNAKE_CASE :			List[str]												=       inspect.signature(model.forward						)
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        __SCREAMING_SNAKE_CASE :			Dict												=       [*signature.parameters.keys()]

        __SCREAMING_SNAKE_CASE :			Optional[Any]												=       ['''pixel_values''']
        self.assertListEqual(arg_names[:1]						,							_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :List[Any]						):
       __SCREAMING_SNAKE_CASE :			Tuple												=       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :List[str]						):
       __SCREAMING_SNAKE_CASE :			int												=       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_segmentation(*_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :int						):
       if not self.model_tester.is_training:
        return

       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			List[Any]												=       self.model_tester.prepare_config_and_inputs_for_common()
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       True

       for model_class in self.all_model_classes:
        if model_class in [*get_values(_lowerCamelCase						)]:
         continue

        __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       model_class(_lowerCamelCase						)
        model.to(_lowerCamelCase						)
        model.train()
        __SCREAMING_SNAKE_CASE :			Optional[int]												=       self._prepare_for_class(_lowerCamelCase						,							_lowerCamelCase						,							return_labels=_lowerCamelCase						)
        __SCREAMING_SNAKE_CASE :			List[Any]												=       model(**_lowerCamelCase						).loss
        loss.backward()

      def 						SCREAMING_SNAKE_CASE_			(			self   :Tuple						):
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			Optional[Any]												=       self.model_tester.prepare_config_and_inputs_for_common()
       if not self.model_tester.is_training:
        return

       __SCREAMING_SNAKE_CASE :			List[Any]												=       False
       __SCREAMING_SNAKE_CASE :			str												=       True

       for model_class in self.all_model_classes:
        if model_class in [*get_values(_lowerCamelCase						)] or not model_class.supports_gradient_checkpointing:
         continue
        # TODO: remove the following 3 lines once we have a MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING
        # this can then be incorporated into _prepare_for_class in test_modeling_common.py
        elif model_class.__name__ == "Data2VecVisionForSemanticSegmentation":
         __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			Optional[int]												=       inputs_dict['''pixel_values'''].shape
         __SCREAMING_SNAKE_CASE :			Dict												=       torch.zeros(
             [self.model_tester.batch_size, height, width]						,							device=_lowerCamelCase						).long()
        __SCREAMING_SNAKE_CASE :			int												=       model_class(_lowerCamelCase						)
        model.gradient_checkpointing_enable()
        model.to(_lowerCamelCase						)
        model.train()
        __SCREAMING_SNAKE_CASE :			Any												=       self._prepare_for_class(_lowerCamelCase						,							_lowerCamelCase						,							return_labels=_lowerCamelCase						)
        __SCREAMING_SNAKE_CASE :			Tuple												=       model(**_lowerCamelCase						).loss
        loss.backward()

      def 						SCREAMING_SNAKE_CASE_			(			self   :int						):
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			List[str]												=       self.model_tester.prepare_config_and_inputs_for_common()

       __SCREAMING_SNAKE_CASE :			Tuple												=       _config_zero_init(_lowerCamelCase						)
       for model_class in self.all_model_classes:
        __SCREAMING_SNAKE_CASE :			str												=       model_class(config=_lowerCamelCase						)
        for name, param in model.named_parameters():
         # we skip lambda parameters as these require special initial values
         # determined by config.layer_scale_init_value
         if "lambda" in name:
          continue
         if param.requires_grad:
          self.assertIn(
              ((param.data.mean() * 1e9).round() / 1e9).item()						,							[0.0, 1.0]						,							msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''						,							)

      def 						SCREAMING_SNAKE_CASE_			(			self   :Tuple						,							_lowerCamelCase   :Dict						,							_lowerCamelCase   :Optional[int]						,							_lowerCamelCase   :Tuple						,							_lowerCamelCase   :List[Any]=2e-4						,							_lowerCamelCase   :str="outputs"						,							_lowerCamelCase   :str=None						):
       # We override with a slightly higher tol value, as semseg models tend to diverge a bit more
       super().check_pt_tf_outputs(_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :List[str]						):
       __SCREAMING_SNAKE_CASE :			str												=       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase						)



      @slow
      def 						SCREAMING_SNAKE_CASE_			(			self   :List[Any]						):
       for model_name in DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __SCREAMING_SNAKE_CASE :			int												=       DataaVecVisionModel.from_pretrained(_lowerCamelCase						)
        self.assertIsNotNone(_lowerCamelCase						)







def 		lowerCAmelCase_	(  ):

 '''simple docstring'''






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






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

      @cached_property
      def 						SCREAMING_SNAKE_CASE_			(			self   :Any						):
       return (
           BeitImageProcessor.from_pretrained('''facebook/data2vec-vision-base-ft1k'''						) if is_vision_available() else None
       )



      @slow
      def 						SCREAMING_SNAKE_CASE_			(			self   :List[str]						):
       __SCREAMING_SNAKE_CASE :			List[Any]												=       DataaVecVisionForImageClassification.from_pretrained('''facebook/data2vec-vision-base-ft1k'''						).to(
           _lowerCamelCase						)

       __SCREAMING_SNAKE_CASE :			List[str]												=       self.default_image_processor
       __SCREAMING_SNAKE_CASE :			List[Any]												=       prepare_img()
       __SCREAMING_SNAKE_CASE :			Optional[Any]												=       image_processor(images=_lowerCamelCase						,							return_tensors='''pt'''						).to(_lowerCamelCase						)

       # forward pass
       with torch.no_grad():
        __SCREAMING_SNAKE_CASE :			int												=       model(**_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       outputs.logits

       # verify the logits
       __SCREAMING_SNAKE_CASE :			List[Any]												=       torch.Size((1, 1_0_0_0)						)
       self.assertEqual(logits.shape						,							_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE :			int												=       torch.tensor([0.3_2_7_7, -0.1_3_9_5, 0.0_9_1_1]						).to(_lowerCamelCase						)

       self.assertTrue(torch.allclose(logits[0, :3]						,							_lowerCamelCase						,							atol=1e-4						)						)

       __SCREAMING_SNAKE_CASE :			Any												=       [model.config.labelaid[i] for i in ['''remote control, remote''', '''tabby, tabby cat''']]
       self.assertEqual(logits[0].topk(2						).indices.cpu().tolist()						,							_lowerCamelCase						)





'''simple docstring'''






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    				=							None

if is_ray_available():
		import ray  # noqa: E402 # isort:skip
		from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
		__UpperCAmelCase    				=							None
		__UpperCAmelCase    				=							None
		__UpperCAmelCase    				=							None


def 					_snake_case   (     A       )						->    Any:
							if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
														lowerCAmelCase__          =	unittest.skip('''test requires Datasets, Faiss, psutil'''       )(A       )
							return test_case






@require_distributed_retrieval
class 		a__				(							a__  ):





							'''simple docstring'''






							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Union[str, 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   )       ->							DPRQuestionEncoderTokenizer:
														return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname					,    '''dpr_tokenizer'''   )   )







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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[int]:
														shutil.rmtree(self.tmpdirname   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														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					,    lowerCamelCase_					,    lowerCamelCase_=1_23_45   )       ->							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					,    lowerCamelCase_   )       ->							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					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_=1_23_45   )       ->							Union[str, 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					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Union[str, Any]:
														# 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					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							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   )       ->							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_pytorch_distributed_retriever(init_retrieval=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							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   )       ->							Dict:
														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   )       ->							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_ray_distributed_retriever(init_retrieval=lowerCamelCase_   )					,    lowerCamelCase_					,    lowerCamelCase_   )
														ray.shutdown()







							@require_ray
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														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   )       ->							List[str]:
														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()