


'''simple docstring'''






class 		a__				:





							'''simple docstring'''






							def __init__(		self					,    lowerCamelCase_   )       ->							List[Any]:
														lowerCAmelCase__          =	data
														lowerCAmelCase__          =	None
														lowerCAmelCase__          =	None







							def __str__(		self   )       ->							List[str]:
														return F"""{self.data}"""




class 		a__				:





							'''simple docstring'''






							def __init__(		self   )       ->							Any:
														lowerCAmelCase__          =	None
														lowerCAmelCase__          =	None







							def __iter__(		self   )       ->							Any:
														lowerCAmelCase__          =	self.head
														while node:
																					yield node.data
																					lowerCAmelCase__          =	node.next







							def __str__(		self   )       ->							Any:
														return "->".join([str(lowerCamelCase_   ) for item in self]   )







							def __len__(		self   )       ->							List[Any]:
														return sum(1 for _ in self   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							List[Any]:
														self.insert_at_nth(0					,    lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							int:
														self.insert_at_nth(len(self   )					,    lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Union[str, Any]:
														lowerCAmelCase__          =	len(self   )

														if not 0 <= index <= length:
																					raise IndexError('''list index out of range'''   )
														lowerCAmelCase__          =	Node(lowerCamelCase_   )
														if self.head is None:
																					lowerCAmelCase__          =	lowerCAmelCase__          =	new_node
														elif index == 0:
																					lowerCAmelCase__          =	new_node
																					lowerCAmelCase__          =	self.head
																					lowerCAmelCase__          =	new_node
														elif index == length:
																					lowerCAmelCase__          =	new_node
																					lowerCAmelCase__          =	self.tail
																					lowerCAmelCase__          =	new_node
														else:
																					lowerCAmelCase__          =	self.head
																					for _ in range(0					,    lowerCamelCase_   ):
																												lowerCAmelCase__          =	temp.next
																					lowerCAmelCase__          =	new_node
																					lowerCAmelCase__          =	temp.previous
																					lowerCAmelCase__          =	temp
																					lowerCAmelCase__          =	new_node







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[Any]:
														return self.delete_at_nth(0   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[Any]:
														return self.delete_at_nth(len(self   ) - 1   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							List[str]:
														lowerCAmelCase__          =	len(self   )

														if not 0 <= index <= length - 1:
																					raise IndexError('''list index out of range'''   )
														lowerCAmelCase__          =	self.head  # default first node
														if length == 1:
																					lowerCAmelCase__          =	lowerCAmelCase__          =	None
														elif index == 0:
																					lowerCAmelCase__          =	self.head.next
																					lowerCAmelCase__          =	None
														elif index == length - 1:
																					lowerCAmelCase__          =	self.tail
																					lowerCAmelCase__          =	self.tail.previous
																					lowerCAmelCase__          =	None
														else:
																					lowerCAmelCase__          =	self.head
																					for _ in range(0					,    lowerCamelCase_   ):
																												lowerCAmelCase__          =	temp.next
																					lowerCAmelCase__          =	temp
																					lowerCAmelCase__          =	temp.previous
																					lowerCAmelCase__          =	temp.next
														return delete_node.data







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							str:
														lowerCAmelCase__          =	self.head

														while current.data != data:  # Find the position to delete
																					if current.next:
																												lowerCAmelCase__          =	current.next
																					else:  # We have reached the end an no value matches
																												raise ValueError('''No data matching given value'''   )

														if current == self.head:
																					self.delete_head()

														elif current == self.tail:
																					self.delete_tail()

														else:  # Before: 1 <--> 2(current) <--> 3
																					lowerCAmelCase__          =	current.next  # 1 --> 3
																					lowerCAmelCase__          =	current.previous  # 1 <--> 3
														return data







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Any:
														return len(self   ) == 0


def 					_snake_case   (     )						->    None:
							lowerCAmelCase__          =	DoublyLinkedList()
							assert linked_list.is_empty() is True
							assert str(A       ) == ""

							try:
														linked_list.delete_head()
														raise AssertionError  # This should not happen.
							except IndexError:
														assert True  # This should happen.

							try:
														linked_list.delete_tail()
														raise AssertionError  # This should not happen.
							except IndexError:
														assert True  # This should happen.

							for i in range(10       ):
														assert len(A       ) == i
														linked_list.insert_at_nth(A ,    i + 1       )
							assert str(A       ) == "->".join(str(A       ) for i in range(1 ,    11       )       )

							linked_list.insert_at_head(0       )
							linked_list.insert_at_tail(11       )
							assert str(A       ) == "->".join(str(A       ) for i in range(0 ,    12       )       )

							assert linked_list.delete_head() == 0
							assert linked_list.delete_at_nth(9       ) == 10
							assert linked_list.delete_tail() == 11
							assert len(A       ) == 9
							assert str(A       ) == "->".join(str(A       ) for i in range(1 ,    10       )       )


if __name__ == "__main__":
		from doctest import testmod

		testmod()



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