



"""simple docstring"""




import csv
import json
import logging
import os
import re
import subprocess
import tempfile
import unittest
import zipfile
from pathlib import Path
from typing import Optional
from unittest import mock

import torch

# We use TF to parse the logs
from accelerate import Accelerator
from accelerate.test_utils.testing import (
    MockingTestCase,
    TempDirTestCase,
    require_comet_ml,
    require_tensorboard,
    require_wandb,
    skip,
)
from accelerate.tracking import CometMLTracker, GeneralTracker
from accelerate.utils import ProjectConfiguration, is_comet_ml_available, is_tensorboard_available


if is_comet_ml_available():
     from comet_ml import OfflineExperiment

if is_tensorboard_available():
     import struct

     import tensorboard.compat.proto.event_pba as event_pba

_UpperCamelCase:					List[str]     				=	logging.getLogger(__name__)


@require_tensorboard
class 	a__		(    unittest.TestCase  ):


   def 							lowercase					(				self   :				List[str]      )     ->							Tuple:
         lowercase	:   Any											= 'test_project_with_config'
         with tempfile.TemporaryDirectory() as dirpath:
               lowercase	:   Optional[Any]											= Accelerator(log_with='tensorboard',  project_dir=_lowerCAmelCase      )
               lowercase	:   Dict											= {'num_iterations': 12, 'learning_rate': 1e-2, 'some_boolean': False, 'some_string': 'some_value'}
               accelerator.init_trackers(_lowerCAmelCase,  _lowerCAmelCase      )
               accelerator.end_training()
               for child in Path(f'''{dirpath}/{project_name}'''      ).glob('*/**'      ):
                     lowercase	:   Dict											= list(filter(lambda lowerCAmelCase      : x.is_file(),  child.iterdir()      )      )[0]
               self.assertNotEqual(str(_lowerCAmelCase      ),  ''      )


   def 							lowercase					(				self   :				List[Any]      )     ->							Union[str, Any]:
         lowercase	:   Dict											= 'test_project_with_log'
         with tempfile.TemporaryDirectory() as dirpath:
               lowercase	:   List[Any]											= Accelerator(log_with='tensorboard',  project_dir=_lowerCAmelCase      )
               accelerator.init_trackers(_lowerCAmelCase      )
               lowercase	:   List[str]											= {'total_loss': 0.1, 'iteration': 1, 'my_text': 'some_value'}
               accelerator.log(_lowerCAmelCase,  step=0      )
               accelerator.end_training()
               # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
               # Names are randomly generated each time
               lowercase	:   int											= list(filter(lambda lowerCAmelCase      : x.is_file(),  Path(f'''{dirpath}/{project_name}'''      ).iterdir()      )      )[0]
               self.assertNotEqual(str(_lowerCAmelCase      ),  ''      )


   def 							lowercase					(				self   :				Tuple      )     ->							List[str]:
         lowercase	:   Dict											= 'test_project_with_log'
         with tempfile.TemporaryDirectory() as dirpath:
               lowercase	:   int											= Accelerator(log_with='tensorboard',  project_dir=_lowerCAmelCase      )
               accelerator.init_trackers(_lowerCAmelCase      )
               lowercase	:   int											= {'tensor': torch.tensor(1      )}
               accelerator.log(_lowerCAmelCase,  step=0      )
               accelerator.end_training()
               # Logged values are stored in the outermost-tfevents file and can be read in as a TFRecord
               # Names are randomly generated each time
               lowercase	:   Any											= list(filter(lambda lowerCAmelCase      : x.is_file(),  Path(f'''{dirpath}/{project_name}'''      ).iterdir()      )      )[0]
               # Reading implementation based on https://github.com/pytorch/pytorch/issues/45327#issuecomment-703757685
               with open(_lowerCAmelCase,  'rb'      ) as f:
                     lowercase	:   List[Any]											= f.read()
               lowercase	:   List[str]											= False
               while data:
                     lowercase	:   Optional[int]											= struct.unpack('Q',  data[:8]      )

                     lowercase	:   List[str]											= data[12 : 12 + int(header[0]      )]  # 8+4
                     lowercase	:   Any											= data[12 + int(header[0]      ) + 4 :]
                     lowercase	:   Optional[Any]											= event_pba.Event()

                     event.ParseFromString(_lowerCAmelCase      )
                     if event.HasField('summary'      ):
                           for value in event.summary.value:
                                 if value.simple_value == 1.0 and value.tag == "tensor":
                                       lowercase	:   Dict											= True
               self.assertTrue(_lowerCAmelCase,  'Converted tensor was not found in the log file!'      )


   def 							lowercase					(				self   :				Tuple      )     ->							int:
         with self.assertRaisesRegex(_lowerCAmelCase,  'Logging with `tensorboard` requires a `logging_dir`'      ):
               lowercase	:   str											= Accelerator(log_with='tensorboard'      )
         with tempfile.TemporaryDirectory() as dirpath:
               lowercase	:   Any											= Accelerator(log_with='tensorboard',  project_dir=_lowerCAmelCase      )






   def 							lowercase					(				self   :				List[str]      )     ->							Optional[int]:
         lowercase	:   Optional[Any]											= ProjectConfiguration(total_limit=30      )
         with tempfile.TemporaryDirectory() as dirpath:
               lowercase	:   Dict											= Accelerator(log_with='tensorboard',  project_dir=_lowerCAmelCase,  project_config=_lowerCAmelCase      )




@require_wandb
@mock.patch.dict(os.environ,	{'WANDB_MODE': 'offline'}  )
class 	a__		(    SCREAMING_SNAKE_CASE__,	SCREAMING_SNAKE_CASE__  ):


   def 							lowercase					(				self   :				Optional[Any]      )     ->							List[str]:
         super().setUp()
         # wandb let's us override where logs are stored to via the WANDB_DIR env var
         self.add_mocks(mock.patch.dict(os.environ,  {'WANDB_DIR': self.tmpdir}      )      )


   @staticmethod
   def 							lowercase					(				lowerCAmelCase   :				str,  lowerCAmelCase   :				str,  lowerCAmelCase   :				bool = True      )     ->							Optional[int]:
         lowercase	:   Any											= Rf'''{section} ([\S\s]*?)\n\n'''
         if record:
               lowercase	:   Union[str, Any]											= Rf'''Record: {pattern}'''
         lowercase	:   Tuple											= re.findall(_lowerCAmelCase,  _lowerCAmelCase      )[0]
         # A config
         if section == "config" or section == "history":
               lowercase	:   List[Any]											= re.findall(R'"([a-zA-Z0-9_.,]+)',  _lowerCAmelCase      )
               return {key: val for key, val in zip(cleaned_record[0::2],  cleaned_record[1::2]      )}
         # Everything else
         else:
               return dict(re.findall(R'(\w+): "([^\s]+)"',  _lowerCAmelCase      )      )






   @skip
   def 							lowercase					(				self   :				List[Any]      )     ->							str:
         lowercase	:   str											= 'test_project_with_config'
         lowercase	:   Tuple											= Accelerator(log_with='wandb'      )
         lowercase	:   Optional[int]											= {'num_iterations': 12, 'learning_rate': 1e-2, 'some_boolean': False, 'some_string': 'some_value'}
         lowercase	:   List[Any]											= {'wandb': {'tags': ['my_tag']}}
         accelerator.init_trackers(_lowerCAmelCase,  _lowerCAmelCase,  _lowerCAmelCase      )
         lowercase	:   Tuple											= {'total_loss': 0.1, 'iteration': 1, 'my_text': 'some_value'}
         accelerator.log(_lowerCAmelCase,  step=0      )
         accelerator.end_training()
         # The latest offline log is stored at wandb/latest-run/*.wandb
         for child in Path(f'''{self.tmpdir}/wandb/latest-run'''      ).glob('*'      ):
               if child.is_file() and child.suffix == ".wandb":
                     lowercase	:   List[Any]											= subprocess.check_output(
                         ['wandb', 'sync', '--view', '--verbose', str(_lowerCAmelCase      )],  env=os.environ.copy()      ).decode('utf8',  'ignore'      )
                     break

        # Check HPS through careful parsing and cleaning
         lowercase	:   str											= self.parse_log(_lowerCAmelCase,  'config'      )
         self.assertEqual(logged_items['num_iterations'],  '12'      )
         self.assertEqual(logged_items['learning_rate'],  '0.01'      )
         self.assertEqual(logged_items['some_boolean'],  'false'      )
         self.assertEqual(logged_items['some_string'],  'some_value'      )
         self.assertEqual(logged_items['some_string'],  'some_value'      )

         # Run tags
         lowercase	:   Optional[int]											= self.parse_log(_lowerCAmelCase,  'run',  _lowerCAmelCase      )
         self.assertEqual(logged_items['tags'],  'my_tag'      )

         # Actual logging
         lowercase	:   Optional[Any]											= self.parse_log(_lowerCAmelCase,  'history'      )
         self.assertEqual(logged_items['total_loss'],  '0.1'      )
         self.assertEqual(logged_items['iteration'],  '1'      )
         self.assertEqual(logged_items['my_text'],  'some_value'      )
         self.assertEqual(logged_items['_step'],  '0'      )




def 	lowercase__  (      self						,   _UpperCAmelCase						,   _UpperCAmelCase					)  ->      Tuple:






      '''simple docstring'''


      lowercase	:   List[str]											= run_name
      lowercase	:   Tuple											= OfflineExperiment(project_name=__SCREAMING_SNAKE_CASE						,   offline_directory=__SCREAMING_SNAKE_CASE					)
      logger.info(f'''Initialized offline CometML project {self.run_name}'''					)
      logger.info('Make sure to log any initial configurations with `self.store_init_configuration` before training!'					)


@require_comet_ml
@mock.patch.object(SCREAMING_SNAKE_CASE__,	'__init__',	SCREAMING_SNAKE_CASE__  )
class 	a__		(    unittest.TestCase  ):


   @staticmethod
   def 							lowercase					(				lowerCAmelCase   :				Union[str, Any],  lowerCAmelCase   :				str,  lowerCAmelCase   :				bool = False      )     ->							List[Any]:
         for log in log_list:
               lowercase	:   Optional[Any]											= json.loads(_lowerCAmelCase      )['payload']
               if is_param and "param" in j.keys():
                     if j["param"]["paramName"] == key:
                           return j["param"]["paramValue"]
               if "log_other" in j.keys():
                     if j["log_other"]["key"] == key:
                           return j["log_other"]["val"]
               if "metric" in j.keys():
                     if j["metric"]["metricName"] == key:
                           return j["metric"]["metricValue"]


   def 							lowercase					(				self   :				List[str]      )     ->							Union[str, Any]:
         with tempfile.TemporaryDirectory() as d:
               lowercase	:   Optional[Any]											= CometMLTracker('test_project_with_config',  _lowerCAmelCase      )
               lowercase	:   str											= Accelerator(log_with=_lowerCAmelCase      )
               lowercase	:   Union[str, Any]											= {'num_iterations': 12, 'learning_rate': 1e-2, 'some_boolean': False, 'some_string': 'some_value'}
               accelerator.init_trackers(_lowerCAmelCase,  _lowerCAmelCase      )
               accelerator.end_training()
               lowercase	:   str											= os.listdir(_lowerCAmelCase      )[0]  # Comet is nice, it's just a zip file here
               # We parse the raw logs
               lowercase	:   str											= os.path.join(_lowerCAmelCase,  _lowerCAmelCase      )
               lowercase	:   Tuple											= zipfile.ZipFile(_lowerCAmelCase,  'r'      )
               lowercase	:   Optional[int]											= archive.open('messages.json'      ).read().decode('utf-8'      )
         lowercase	:   int											= log.split('\n'      )[:-1]
         self.assertEqual(self.get_value_from_key(_lowerCAmelCase,  'num_iterations',  _lowerCAmelCase      ),  12      )
         self.assertEqual(self.get_value_from_key(_lowerCAmelCase,  'learning_rate',  _lowerCAmelCase      ),  0.01      )
         self.assertEqual(self.get_value_from_key(_lowerCAmelCase,  'some_boolean',  _lowerCAmelCase      ),  _lowerCAmelCase      )
         self.assertEqual(self.get_value_from_key(_lowerCAmelCase,  'some_string',  _lowerCAmelCase      ),  'some_value'      )






   def 							lowercase					(				self   :				List[str]      )     ->							Optional[Any]:
         with tempfile.TemporaryDirectory() as d:
               lowercase	:   Any											= CometMLTracker('test_project_with_config',  _lowerCAmelCase      )
               lowercase	:   Optional[int]											= Accelerator(log_with=_lowerCAmelCase      )
               accelerator.init_trackers(_lowerCAmelCase      )
               lowercase	:   List[str]											= {'total_loss': 0.1, 'iteration': 1, 'my_text': 'some_value'}
               accelerator.log(_lowerCAmelCase,  step=0      )
               accelerator.end_training()
               lowercase	:   Optional[Any]											= os.listdir(_lowerCAmelCase      )[0]  # Comet is nice, it's just a zip file here
               # We parse the raw logs
               lowercase	:   Optional[Any]											= os.path.join(_lowerCAmelCase,  _lowerCAmelCase      )
               lowercase	:   str											= zipfile.ZipFile(_lowerCAmelCase,  'r'      )
               lowercase	:   Tuple											= archive.open('messages.json'      ).read().decode('utf-8'      )
         lowercase	:   Any											= log.split('\n'      )[:-1]
         self.assertEqual(self.get_value_from_key(_lowerCAmelCase,  'curr_step',  _lowerCAmelCase      ),  0      )
         self.assertEqual(self.get_value_from_key(_lowerCAmelCase,  'total_loss'      ),  0.1      )
         self.assertEqual(self.get_value_from_key(_lowerCAmelCase,  'iteration'      ),  1      )
         self.assertEqual(self.get_value_from_key(_lowerCAmelCase,  'my_text'      ),  'some_value'      )




class 	a__		(    SCREAMING_SNAKE_CASE__  ):
   _lowerCamelCase       				=					[
       'total_loss',
       'iteration',
       'my_text',
       'learning_rate',
       'num_iterations',
       'some_boolean',
       'some_string',
   ]

   _lowerCamelCase       				=					'my_custom_tracker'
   _lowerCamelCase       				=					False


   def __init__(				self   :				int,  lowerCAmelCase   :				str      )     ->							Optional[Any]:
         lowercase	:   Optional[int]											= open(f'''{dir}/log.csv''',  'w+'      )
         lowercase	:   Optional[Any]											= csv.DictWriter(self.f,  fieldnames=self._col_names      )
         self.writer.writeheader()


   @property
   def 							lowercase					(				self   :				Any      )     ->							List[str]:
         return self.writer


   def 							lowercase					(				self   :				Optional[int],  lowerCAmelCase   :				dict      )     ->							Optional[Any]:
         logger.info('Call init'      )
         self.writer.writerow(_lowerCAmelCase      )


   def 							lowercase					(				self   :				int,  lowerCAmelCase   :				dict,  lowerCAmelCase   :				Optional[int]      )     ->							Tuple:
         logger.info('Call log'      )
         self.writer.writerow(_lowerCAmelCase      )






   def 							lowercase					(				self   :				Optional[int]      )     ->							Optional[int]:
         self.f.close()




class 	a__		(    unittest.TestCase  ):


   def 							lowercase					(				self   :				List[str]      )     ->							Optional[Any]:
         with tempfile.TemporaryDirectory() as d:
               lowercase	:   Dict											= MyCustomTracker(_lowerCAmelCase      )
               lowercase	:   str											= Accelerator(log_with=_lowerCAmelCase      )
               lowercase	:   List[str]											= {'num_iterations': 12, 'learning_rate': 1e-2, 'some_boolean': False, 'some_string': 'some_value'}
               accelerator.init_trackers('Some name',  _lowerCAmelCase      )
               accelerator.end_training()
               with open(f'''{d}/log.csv''',  'r'      ) as f:
                     lowercase	:   str											= csv.DictReader(_lowerCAmelCase      )
                     lowercase	:   Optional[int]											= next(_lowerCAmelCase      )
                     lowercase	:   int											= {
                         'total_loss': '',
                         'iteration': '',
                         'my_text': '',
                         'learning_rate': '0.01',
                         'num_iterations': '12',
                         'some_boolean': 'False',
                         'some_string': 'some_value',
                     }
                     self.assertDictEqual(_lowerCAmelCase,  _lowerCAmelCase      )






   def 							lowercase					(				self   :				Optional[Any]      )     ->							Optional[Any]:
         with tempfile.TemporaryDirectory() as d:
               lowercase	:   List[Any]											= MyCustomTracker(_lowerCAmelCase      )
               lowercase	:   List[Any]											= Accelerator(log_with=_lowerCAmelCase      )
               accelerator.init_trackers('Some name'      )
               lowercase	:   List[Any]											= {'total_loss': 0.1, 'iteration': 1, 'my_text': 'some_value'}
               accelerator.log(_lowerCAmelCase,  step=0      )
               accelerator.end_training()
               with open(f'''{d}/log.csv''',  'r'      ) as f:
                     lowercase	:   Union[str, Any]											= csv.DictReader(_lowerCAmelCase      )
                     lowercase	:   str											= next(_lowerCAmelCase      )
                     lowercase	:   List[Any]											= {
                         'total_loss': '0.1',
                         'iteration': '1',
                         'my_text': 'some_value',
                         'learning_rate': '',
                         'num_iterations': '',
                         'some_boolean': '',
                         'some_string': '',
                     }
                     self.assertDictEqual(_lowerCAmelCase,  _lowerCAmelCase      )


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()
