
"""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:
					_A            = None

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





		'''simple docstring'''


		if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
				lowerCamelCase     :			Optional[int]		=    unittest.skip('test requires Datasets, Faiss, psutil'			)(a_			)
		return test_case






@require_distributed_retrieval
class        _lowercase   (				__UpperCAmelCase		):




	def 			_UpperCamelCase    ( self					)		->							Any:
			lowerCamelCase     :			Tuple		=    tempfile.mkdtemp()
			lowerCamelCase     :			Any		=    8

			# DPR tok
			lowerCamelCase     :			str		=    [
			    '[UNK]',
			    '[CLS]',
			    '[SEP]',
			    '[PAD]',
			    '[MASK]',
			    'want',
			    '##want',
			    '##ed',
			    'wa',
			    'un',
			    'runn',
			    '##ing',
			    ',',
			    'low',
			    'lowest',
			]
			lowerCamelCase     :			Optional[Any]		=    os.path.join(self.tmpdirname		, 'dpr_tokenizer'					)
			os.makedirs(UpperCAmelCase_		, exist_ok=UpperCAmelCase_					)
			lowerCamelCase     :			int		=    os.path.join(UpperCAmelCase_		, 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     :			List[Any]		=    [
			    'l',
			    'o',
			    'w',
			    'e',
			    'r',
			    's',
			    't',
			    'i',
			    'd',
			    'n',
			    '\u0120',
			    '\u0120l',
			    '\u0120n',
			    '\u0120lo',
			    '\u0120low',
			    'er',
			    '\u0120lowest',
			    '\u0120newer',
			    '\u0120wider',
			    '<unk>',
			]
			lowerCamelCase     :			Tuple		=    dict(zip(UpperCAmelCase_		, range(len(UpperCAmelCase_					)					)					)					)
			lowerCamelCase     :			Tuple		=    ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
			lowerCamelCase     :			Union[str, Any]		=    {'unk_token': '<unk>'}

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




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




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




	def 			_UpperCamelCase    ( self					)		->							Any:
			shutil.rmtree(self.tmpdirname					)




	def 			_UpperCamelCase    ( self					)		->							str:
			lowerCamelCase     :			Optional[int]		=    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 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_=12345					)		->							RagPyTorchDistributedRetriever:
			lowerCamelCase     :			List[Any]		=    self.get_dummy_dataset()
			lowerCamelCase     :			Tuple		=    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     :			Optional[int]		=    dataset
					lowerCamelCase     :			int		=    RagPyTorchDistributedRetriever(
					    UpperCAmelCase_		, question_encoder_tokenizer=self.get_dpr_tokenizer()		, generator_tokenizer=self.get_bart_tokenizer()		, )
					if init_retrieval:
							retriever.init_retrieval(UpperCAmelCase_					)
			return retriever




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_					)		->							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=UpperCAmelCase_					)
			lowerCamelCase     :			List[Any]		=    RagConfig(
			    retrieval_vector_size=self.retrieval_vector_size		, question_encoder=DPRConfig().to_dict()		, generator=BartConfig().to_dict()		, )
			lowerCamelCase     :			Optional[Any]		=    ray.remote(UpperCAmelCase_					)
			lowerCamelCase     :			Optional[Any]		=    [remote_cls.remote() for _ in range(1					)]
			with patch('transformers.models.rag.retrieval_rag.load_dataset'					) as mock_load_dataset:
					lowerCamelCase     :			Dict		=    self.get_dummy_dataset()
					lowerCamelCase     :			Union[str, Any]		=    RagRayDistributedRetriever(
					    UpperCAmelCase_		, question_encoder_tokenizer=self.get_dpr_tokenizer()		, generator_tokenizer=self.get_bart_tokenizer()		, retrieval_workers=UpperCAmelCase_		, )
					if init_retrieval:
							retriever.init_retrieval()
			return retriever




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_=12345					)		->							Tuple:
			lowerCamelCase     :			Optional[Any]		=    self.get_dummy_dataset()
			lowerCamelCase     :			Union[str, Any]		=    RagConfig(
			    retrieval_vector_size=self.retrieval_vector_size		, question_encoder=DPRConfig().to_dict()		, generator=BartConfig().to_dict()		, index_name='custom'		, )
			if from_disk:
					lowerCamelCase     :			Tuple		=    os.path.join(self.tmpdirname		, 'dataset'					)
					lowerCamelCase     :			str		=    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     :			List[str]		=    RagPyTorchDistributedRetriever(
					    UpperCAmelCase_		, question_encoder_tokenizer=self.get_dpr_tokenizer()		, generator_tokenizer=self.get_bart_tokenizer()		, )
			else:
					lowerCamelCase     :			List[Any]		=    RagPyTorchDistributedRetriever(
					    UpperCAmelCase_		, question_encoder_tokenizer=self.get_dpr_tokenizer()		, generator_tokenizer=self.get_bart_tokenizer()		, index=CustomHFIndex(config.retrieval_vector_size		, UpperCAmelCase_					)		, )
			if init_retrieval:
					retriever.init_retrieval(UpperCAmelCase_					)
			return retriever




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_					)		->							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=UpperCAmelCase_					)
			lowerCamelCase     :			List[Any]		=    self.get_dummy_dataset()
			lowerCamelCase     :			List[Any]		=    RagConfig(
			    retrieval_vector_size=self.retrieval_vector_size		, question_encoder=DPRConfig().to_dict()		, generator=BartConfig().to_dict()		, index_name='custom'		, )
			lowerCamelCase     :			Optional[Any]		=    ray.remote(UpperCAmelCase_					)
			lowerCamelCase     :			Optional[int]		=    [remote_cls.remote() for _ in range(1					)]
			if from_disk:
					lowerCamelCase     :			List[Any]		=    os.path.join(self.tmpdirname		, 'dataset'					)
					lowerCamelCase     :			Optional[Any]		=    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     :			List[str]		=    RagRayDistributedRetriever(
					    UpperCAmelCase_		, question_encoder_tokenizer=self.get_dpr_tokenizer()		, generator_tokenizer=self.get_bart_tokenizer()		, retrieval_workers=UpperCAmelCase_		, index=CustomHFIndex.load_from_disk(
					        vector_size=config.retrieval_vector_size		, dataset_path=config.passages_path		, index_path=config.index_path		, )		, )
			else:
					lowerCamelCase     :			Any		=    RagRayDistributedRetriever(
					    UpperCAmelCase_		, question_encoder_tokenizer=self.get_dpr_tokenizer()		, generator_tokenizer=self.get_bart_tokenizer()		, retrieval_workers=UpperCAmelCase_		, index=CustomHFIndex(config.retrieval_vector_size		, UpperCAmelCase_					)		, )
			if init_retrieval:
					retriever.init_retrieval()
			return retriever




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)		->							None:
			lowerCamelCase	,    lowerCamelCase	,    lowerCamelCase     :			Any		=    retriever.retrieve(UpperCAmelCase_		, n_docs=UpperCAmelCase_					)
			self.assertEqual(retrieved_doc_embeds.shape		, (2, n_docs, self.retrieval_vector_size)					)
			self.assertEqual(len(UpperCAmelCase_					)		, 2					)
			self.assertEqual(sorted(doc_dicts[0]					)		, ['embeddings', 'id', 'text', 'title']					)
			self.assertEqual(len(doc_dicts[0]['id']					)		, UpperCAmelCase_					)
			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 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			lowerCamelCase     :			List[str]		=    1
			lowerCamelCase     :			Dict		=    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=UpperCAmelCase_					)		, UpperCAmelCase_		, UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							List[str]:
			lowerCamelCase     :			Optional[int]		=    1
			lowerCamelCase     :			Union[str, Any]		=    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=UpperCAmelCase_		, from_disk=UpperCAmelCase_					)		, UpperCAmelCase_		, UpperCAmelCase_		, )




	def 			_UpperCamelCase    ( self					)		->							str:
			lowerCamelCase     :			List[Any]		=    1
			lowerCamelCase     :			Optional[int]		=    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=UpperCAmelCase_		, from_disk=UpperCAmelCase_					)		, UpperCAmelCase_		, UpperCAmelCase_		, )




	@require_ray
	def 			_UpperCamelCase    ( self					)		->							Tuple:
			lowerCamelCase     :			int		=    1
			lowerCamelCase     :			Union[str, Any]		=    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=UpperCAmelCase_					)		, UpperCAmelCase_		, UpperCAmelCase_					)
			ray.shutdown()




	@require_ray
	def 			_UpperCamelCase    ( self					)		->							int:
			lowerCamelCase     :			Optional[int]		=    1
			lowerCamelCase     :			Optional[int]		=    np.array(
			    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]		, dtype=np.floataa					)
			with self.assertRaises(UpperCAmelCase_					):
					self.distributed_retriever_check(
					    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=UpperCAmelCase_		, from_disk=UpperCAmelCase_					)		, UpperCAmelCase_		, UpperCAmelCase_		, )
			ray.shutdown()




	@require_ray
	def 			_UpperCamelCase    ( self					)		->							Dict:
			lowerCamelCase     :			str		=    1
			lowerCamelCase     :			List[Any]		=    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=UpperCAmelCase_		, from_disk=UpperCAmelCase_					)		, UpperCAmelCase_		, UpperCAmelCase_					)
			ray.shutdown()








"""simple docstring"""

from transformers import (
    DPRContextEncoderTokenizer,
    DPRContextEncoderTokenizerFast,
    DPRQuestionEncoderTokenizer,
    DPRQuestionEncoderTokenizerFast,
    DPRReaderOutput,
    DPRReaderTokenizer,
    DPRReaderTokenizerFast,
)
from transformers.testing_utils import require_tokenizers, slow
from transformers.tokenization_utils_base import BatchEncoding

from ..bert.test_tokenization_bert import BertTokenizationTest






@require_tokenizers
class        _lowercase   (				__UpperCAmelCase		):
	lowercase_        =	DPRContextEncoderTokenizer
	lowercase_        =	DPRContextEncoderTokenizerFast
	lowercase_        =	True






@require_tokenizers
class        _lowercase   (				__UpperCAmelCase		):
	lowercase_        =	DPRQuestionEncoderTokenizer
	lowercase_        =	DPRQuestionEncoderTokenizerFast
	lowercase_        =	True






@require_tokenizers
class        _lowercase   (				__UpperCAmelCase		):
	lowercase_        =	DPRReaderTokenizer
	lowercase_        =	DPRReaderTokenizerFast
	lowercase_        =	True




	@slow
	def 			_UpperCamelCase    ( self					)		->							Any:
			lowerCamelCase     :			str		=    self.tokenizer_class.from_pretrained('bert-base-uncased'					)

			lowerCamelCase     :			int		=    tokenizer.encode('question sequence'		, add_special_tokens=UpperCAmelCase_					)
			lowerCamelCase     :			List[Any]		=    tokenizer.encode('title sequence'		, add_special_tokens=UpperCAmelCase_					)
			lowerCamelCase     :			Optional[int]		=    tokenizer.encode('text sequence ' * 4		, add_special_tokens=UpperCAmelCase_					)
			lowerCamelCase     :			Union[str, Any]		=    [[101] + text_a + [102] + text_a + [102] + text_a]
			lowerCamelCase     :			List[str]		=    BatchEncoding({'input_ids': input_ids}					)

			lowerCamelCase     :			Optional[int]		=    [[0] * len(input_ids[0]					)]
			lowerCamelCase     :			str		=    [[0] * len(input_ids[0]					)]
			lowerCamelCase     :			List[Any]		=    [0]
			lowerCamelCase     :			Tuple		=    DPRReaderOutput(UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)

			lowerCamelCase	,    lowerCamelCase     :			Optional[int]		=    8, 9
			lowerCamelCase     :			Optional[int]		=    10
			lowerCamelCase     :			int		=    10
			lowerCamelCase     :			List[Any]		=    tokenizer.decode_best_spans(UpperCAmelCase_		, UpperCAmelCase_					)
			self.assertEqual(predicted_spans[0].start_index		, UpperCAmelCase_					)
			self.assertEqual(predicted_spans[0].end_index		, UpperCAmelCase_					)
			self.assertEqual(predicted_spans[0].doc_id		, 0					)




	@slow
	def 			_UpperCamelCase    ( self					)		->							List[Any]:
			lowerCamelCase     :			Union[str, Any]		=    self.tokenizer_class.from_pretrained('bert-base-uncased'					)

			lowerCamelCase     :			List[str]		=    tokenizer.encode('question sequence'		, add_special_tokens=UpperCAmelCase_					)
			lowerCamelCase     :			List[str]		=    tokenizer.encode('title sequence'		, add_special_tokens=UpperCAmelCase_					)
			lowerCamelCase     :			List[Any]		=    tokenizer.encode('text sequence'		, add_special_tokens=UpperCAmelCase_					)
			lowerCamelCase     :			List[Any]		=    [101] + text_a + [102] + text_a + [102] + text_a
			lowerCamelCase     :			Optional[Any]		=    tokenizer(questions=['question sequence']		, titles=['title sequence']		, texts=['text sequence']					)
			self.assertIn('input_ids'		, UpperCAmelCase_					)
			self.assertIn('attention_mask'		, UpperCAmelCase_					)
			self.assertListEqual(encoded_input['input_ids'][0]		, UpperCAmelCase_					)






