




import logging
import random

import ray

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


__A						:		str          =						logging.getLogger(__name__)

class 		lowercase_       :




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


									'''simple docstring'''


									__lowerCAmelCase					=  False




					def     _lowercase						(   self:		str,     _lowercase:		Tuple,     _lowercase:		List[str],     _lowercase:		List[str],     _lowercase:		Union[str, Any]):


									'''simple docstring'''


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




					def     _lowercase						(   self:		str):


									'''simple docstring'''


									self.retriever.index.init_index()




					def     _lowercase						(   self:		List[str]):


									'''simple docstring'''


									del self.retriever
									__lowerCAmelCase					=  False







					def     _lowercase						(   self:		Any,     _lowercase:		Optional[int],     _lowercase:		Dict):


									'''simple docstring'''


									__lowerCAmelCase	,				__lowerCAmelCase					=  self.retriever._main_retrieve(_lowercase,     _lowercase)
									__lowerCAmelCase					=  self.retriever.index.get_doc_dicts(_lowercase)
									return doc_ids, retrieved_doc_embeds, doc_dicts





class 		lowercase_       (	lowerCAmelCase__				):




					def __init__(   self:		str,     _lowercase:		List[Any],     _lowercase:		str,     _lowercase:		Dict,     _lowercase:		List[str],     _lowercase:		str=None):


									'''simple docstring'''


									if index is not None and index.is_initialized() and len(_lowercase) > 0:
													raise ValueError(
													    """When using Ray for distributed fine-tuning, """
													    """you'll need to provide the paths instead, """
													    """as the dataset and the index are loaded """
													    """separately. More info in examples/rag/use_own_knowledge_dataset.py """)

									super().__init__(
									    _lowercase,     question_encoder_tokenizer=_lowercase,     generator_tokenizer=_lowercase,     index=_lowercase,     init_retrieval=_lowercase,     )

									__lowerCAmelCase					=  retrieval_workers
									__lowerCAmelCase					=  question_encoder_tokenizer
									__lowerCAmelCase					=  generator_tokenizer
									if len(self.retrieval_workers) > 0:
													ray.get(
													    [
													        worker.create_rag_retriever.remote(_lowercase,     _lowercase,     _lowercase,     _lowercase)
													        for worker in self.retrieval_workers
													    ])




					def     _lowercase						(   self:		Tuple):


									'''simple docstring'''


									logger.info("""initializing retrieval""")

									if len(self.retrieval_workers) > 0:
													ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers])
									else:
													# Non-distributed training. Load index into this same process.
													self.index.init_index()




					def     _lowercase						(   self:		Optional[Any],     _lowercase:		Optional[Any],     _lowercase:		List[str]):


									'''simple docstring'''


									if len(self.retrieval_workers) > 0:
													# Select a random retrieval actor.
													__lowerCAmelCase					=  self.retrieval_workers[random.randint(0,     len(self.retrieval_workers) - 1)]
													__lowerCAmelCase	,				__lowerCAmelCase	,				__lowerCAmelCase					=  ray.get(
													    random_worker.retrieve.remote(_lowercase,     _lowercase))
									else:
													__lowerCAmelCase	,				__lowerCAmelCase					=  self._main_retrieve(_lowercase,     _lowercase)
													__lowerCAmelCase					=  self.index.get_doc_dicts(_lowercase)
									return retrieved_doc_embeds, doc_ids, doc_dicts




					@classmethod
					def     _lowercase						(   cls:		Union[str, Any],     _lowercase:		Optional[int],     _lowercase:		Any=None,     **_lowercase:		str):


									'''simple docstring'''


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




					@classmethod
					def     _lowercase						(   cls:		Union[str, Any],     _lowercase:		List[Any],     _lowercase:		List[Any],     _lowercase:		List[str]=None,     **_lowercase:		Any):


									'''simple docstring'''


									__lowerCAmelCase					=  kwargs.pop("""config""",     _lowercase) or RagConfig.from_pretrained(_lowercase,     **_lowercase)
									__lowerCAmelCase					=  RagTokenizer.from_pretrained(_lowercase,     config=_lowercase)
									__lowerCAmelCase					=  rag_tokenizer.question_encoder
									__lowerCAmelCase					=  rag_tokenizer.generator

									if indexed_dataset is not None:
													__lowerCAmelCase					=  """custom"""
													__lowerCAmelCase					=  CustomHFIndex(config.retrieval_vector_size,     _lowercase)
									else:
													__lowerCAmelCase					=  cls._build_index(_lowercase)

									return cls(
									    _lowercase,     question_encoder_tokenizer=_lowercase,     generator_tokenizer=_lowercase,     retrieval_workers=_lowercase,     index=_lowercase,     )







					def     _lowercase						(   self:		int):


									'''simple docstring'''


									logger.info("""re-loading the new dataset with embeddings""")
									# access from the training loop

									ray.get([worker.clear_object.remote() for worker in self.retrieval_workers])

									# build the index object again
									__lowerCAmelCase					=  self._build_index(self.config)

									ray.get(
									    [
									        worker.create_rag_retriever.remote(
									            self.config,     self.question_encoder_tokenizer,     self.generator_tokenizer,     _lowercase)
									        for worker in self.retrieval_workers
									    ])











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:
	__A						:		str          =						None

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





def      UpperCAmelCase  (						UpperCamelCase__   )     ->     Union[str, Any]:







				'''simple docstring'''


				if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
								__lowerCAmelCase					=  unittest.skip("""test requires Datasets, Faiss, psutil"""   )(UpperCamelCase__   )
				return test_case

@require_distributed_retrieval
class 		lowercase_       (	lowerCAmelCase__				):




					def     _lowercase						(   self:		List[Any]):


									'''simple docstring'''


									__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(_lowercase,     exist_ok=_lowercase)
									__lowerCAmelCase					=  os.path.join(_lowercase,     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(_lowercase,     range(len(_lowercase))))
									__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(_lowercase,     exist_ok=_lowercase)
									__lowerCAmelCase					=  os.path.join(_lowercase,     BART_VOCAB_FILES_NAMES["""vocab_file"""])
									__lowerCAmelCase					=  os.path.join(_lowercase,     BART_VOCAB_FILES_NAMES["""merges_file"""])
									with open(self.vocab_file,     """w""",     encoding="""utf-8""") as fp:
													fp.write(json.dumps(_lowercase) + """\n""")
									with open(self.merges_file,     """w""",     encoding="""utf-8""") as fp:
													fp.write("""\n""".join(_lowercase))




					def     _lowercase						(   self:		Optional[Any]):


									'''simple docstring'''


									return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname,     """dpr_tokenizer"""))




					def     _lowercase						(   self:		int):


									'''simple docstring'''


									return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname,     """bart_tokenizer"""))




					def     _lowercase						(   self:		Optional[int]):


									'''simple docstring'''


									shutil.rmtree(self.tmpdirname)




					def     _lowercase						(   self:		Tuple):


									'''simple docstring'''


									__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     _lowercase						(   self:		Union[str, Any],     _lowercase:		bool,     _lowercase:		Any=12345):


									'''simple docstring'''


									__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(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     )
													if init_retrieval:
																	retriever.init_retrieval(_lowercase)
									return retriever




					def     _lowercase						(   self:		Union[str, Any],     _lowercase:		bool):


									'''simple docstring'''


									ray.init(local_mode=_lowercase)
									__lowerCAmelCase					=  RagConfig(
									    retrieval_vector_size=self.retrieval_vector_size,     question_encoder=DPRConfig().to_dict(),     generator=BartConfig().to_dict(),     )
									__lowerCAmelCase					=  ray.remote(_lowercase)
									__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(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     )
													if init_retrieval:
																	retriever.init_retrieval()
									return retriever




					def     _lowercase						(   self:		Optional[Any],     _lowercase:		bool,     _lowercase:		bool,     _lowercase:		int=12345):


									'''simple docstring'''


									__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(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     )
									else:
													__lowerCAmelCase					=  RagPyTorchDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     index=CustomHFIndex(config.retrieval_vector_size,     _lowercase),     )
									if init_retrieval:
													retriever.init_retrieval(_lowercase)
									return retriever




					def     _lowercase						(   self:		List[Any],     _lowercase:		bool,     _lowercase:		bool):


									'''simple docstring'''


									ray.init(local_mode=_lowercase)
									__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(_lowercase)
									__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(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     index=CustomHFIndex.load_from_disk(
													        vector_size=config.retrieval_vector_size,     dataset_path=config.passages_path,     index_path=config.index_path,     ),     )
									else:
													__lowerCAmelCase					=  RagRayDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     index=CustomHFIndex(config.retrieval_vector_size,     _lowercase),     )
									if init_retrieval:
													retriever.init_retrieval()
									return retriever




					def     _lowercase						(   self:		int,     _lowercase:		RagRetriever,     _lowercase:		np.array,     _lowercase:		int):


									'''simple docstring'''


									__lowerCAmelCase	,				__lowerCAmelCase	,				__lowerCAmelCase					=  retriever.retrieve(_lowercase,     n_docs=_lowercase)
									self.assertEqual(retrieved_doc_embeds.shape,     (2, n_docs, self.retrieval_vector_size))
									self.assertEqual(len(_lowercase),     2)
									self.assertEqual(sorted(doc_dicts[0]),     ["""embeddings""", """id""", """text""", """title"""])
									self.assertEqual(len(doc_dicts[0]["""id"""]),     _lowercase)
									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     _lowercase						(   self:		Optional[Any]):


									'''simple docstring'''


									__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=_lowercase),     _lowercase,     _lowercase)




					def     _lowercase						(   self:		int):


									'''simple docstring'''


									__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=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )




					def     _lowercase						(   self:		Dict):


									'''simple docstring'''


									__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=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )




					@require_ray
					def     _lowercase						(   self:		List[str]):


									'''simple docstring'''


									__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=_lowercase),     _lowercase,     _lowercase)
									ray.shutdown()




					@require_ray
					def     _lowercase						(   self:		Union[str, Any]):


									'''simple docstring'''


									__lowerCAmelCase					=  1
									__lowerCAmelCase					=  np.array(
									    [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)],     dtype=np.floataa)
									with self.assertRaises(_lowercase):
													self.distributed_retriever_check(
													    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )
									ray.shutdown()







					@require_ray
					def     _lowercase						(   self:		Tuple):


									'''simple docstring'''


									__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=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase)
									ray.shutdown()





