
"""simple docstring"""
from __future__ import annotations

import unittest
from collections.abc import Iterator




class    _snake_case  :






							def __init__(					self						:			List[Any]						, snake_case__						:			int						, snake_case__						:			Node | None    ):






												"""simple docstring"""



												__a:						Tuple		    =					label
												__a:						Union[str, Any]		    =					parent
												__a:						Node | None		    =					None
												__a:						Node | None		    =					None






class    _snake_case  :






							def __init__(					self						:			str    ):






												"""simple docstring"""



												__a:						Node | None		    =					None






							def 							__lowercase  (					self						:			List[Any]    ):






												"""simple docstring"""



												__a:						Dict		    =					None






							def 							__lowercase  (					self						:			Any    ):






												"""simple docstring"""



												return self.root is None






							def 							__lowercase  (					self						:			Tuple						, snake_case__						:			int    ):






												"""simple docstring"""



												__a:						Tuple		    =					self._put(self.root						, snake_case__    )






							def 							__lowercase  (					self						:			List[str]						, snake_case__						:			Node | None						, snake_case__						:			int						, snake_case__						:			Node | None = None    ):






												"""simple docstring"""



												if node is None:
																	__a:						str		    =					Node(snake_case__						, snake_case__    )
												else:
																	if label < node.label:
																						__a:						List[str]		    =					self._put(node.left						, snake_case__						, snake_case__    )
																	elif label > node.label:
																						__a:						str		    =					self._put(node.right						, snake_case__						, snake_case__    )
																	else:
																						__a:						int		    =					f"""Node with label {label} already exists"""
																						raise Exception(snake_case__    )

												return node






							def 							__lowercase  (					self						:			List[Any]						, snake_case__						:			int    ):






												"""simple docstring"""



												return self._search(self.root						, snake_case__    )






							def 							__lowercase  (					self						:			Tuple						, snake_case__						:			Node | None						, snake_case__						:			int    ):






												"""simple docstring"""



												if node is None:
																	__a:						Dict		    =					f"""Node with label {label} does not exist"""
																	raise Exception(snake_case__    )
												else:
																	if label < node.label:
																						__a:						Union[str, Any]		    =					self._search(node.left						, snake_case__    )
																	elif label > node.label:
																						__a:						Optional[int]		    =					self._search(node.right						, snake_case__    )

												return node






							def 							__lowercase  (					self						:			Dict						, snake_case__						:			int    ):






												"""simple docstring"""



												__a:						str		    =					self.search(snake_case__    )
												if node.right and node.left:
																	__a:						Any		    =					self._get_lowest_node(node.right    )
																	__a:						List[Any]		    =					node.left
																	__a:						List[str]		    =					node.right
																	__a:						Union[str, Any]		    =					lowest_node
																	if node.right:
																						__a:						Dict		    =					lowest_node
																	self._reassign_nodes(snake_case__						, snake_case__    )
												elif not node.right and node.left:
																	self._reassign_nodes(snake_case__						, node.left    )
												elif node.right and not node.left:
																	self._reassign_nodes(snake_case__						, node.right    )
												else:
																	self._reassign_nodes(snake_case__						, snake_case__    )






							def 							__lowercase  (					self						:			Optional[int]						, snake_case__						:			Node						, snake_case__						:			Node | None    ):






												"""simple docstring"""



												if new_children:
																	__a:						Optional[Any]		    =					node.parent

												if node.parent:
																	if node.parent.right == node:
																						__a:						List[Any]		    =					new_children
																	else:
																						__a:						int		    =					new_children
												else:
																	__a:						str		    =					new_children






							def 							__lowercase  (					self						:			Tuple						, snake_case__						:			Node    ):






												"""simple docstring"""



												if node.left:
																	__a:						Tuple		    =					self._get_lowest_node(node.left    )
												else:
																	__a:						str		    =					node
																	self._reassign_nodes(snake_case__						, node.right    )

												return lowest_node






							def 							__lowercase  (					self						:			str						, snake_case__						:			int    ):






												"""simple docstring"""



												try:
																	self.search(snake_case__    )
																	return True
												except Exception:
																	return False






							def 							__lowercase  (					self						:			Optional[int]    ):






												"""simple docstring"""



												if self.root is None:
																	raise Exception("""Binary search tree is empty"""    )

												__a:						int		    =					self.root
												while node.right is not None:
																	__a:						List[Any]		    =					node.right

												return node.label






							def 							__lowercase  (					self						:			Optional[Any]    ):






												"""simple docstring"""



												if self.root is None:
																	raise Exception("""Binary search tree is empty"""    )

												__a:						List[str]		    =					self.root
												while node.left is not None:
																	__a:						List[Any]		    =					node.left

												return node.label






							def 							__lowercase  (					self						:			List[str]    ):






												"""simple docstring"""



												return self._inorder_traversal(self.root    )






							def 							__lowercase  (					self						:			Any						, snake_case__						:			Node | None    ):






												"""simple docstring"""



												if node is not None:
																	yield from self._inorder_traversal(node.left    )
																	yield node
																	yield from self._inorder_traversal(node.right    )






							def 							__lowercase  (					self						:			List[str]    ):






												"""simple docstring"""



												return self._preorder_traversal(self.root    )






							def 							__lowercase  (					self						:			Any						, snake_case__						:			Node | None    ):






												"""simple docstring"""



												if node is not None:
																	yield node
																	yield from self._preorder_traversal(node.left    )
																	yield from self._preorder_traversal(node.right    )






class    _snake_case  (						unittest.TestCase       ):






							@staticmethod
							def 							__lowercase  (					):






												"""simple docstring"""



												__a:						List[str]		    =					BinarySearchTree()
												t.put(8    )
												t.put(3    )
												t.put(6    )
												t.put(1    )
												t.put(1_0    )
												t.put(1_4    )
												t.put(1_3    )
												t.put(4    )
												t.put(7    )
												t.put(5    )

												return t






							def 							__lowercase  (					self						:			int    ):






												"""simple docstring"""



												__a:						Union[str, Any]		    =					BinarySearchTree()
												assert t.is_empty()

												t.put(8    )
												assert t.root is not None
												assert t.root.parent is None
												assert t.root.label == 8

												t.put(1_0    )
												assert t.root.right is not None
												assert t.root.right.parent == t.root
												assert t.root.right.label == 1_0

												t.put(3    )
												assert t.root.left is not None
												assert t.root.left.parent == t.root
												assert t.root.left.label == 3

												t.put(6    )
												assert t.root.left.right is not None
												assert t.root.left.right.parent == t.root.left
												assert t.root.left.right.label == 6

												t.put(1    )
												assert t.root.left.left is not None
												assert t.root.left.left.parent == t.root.left
												assert t.root.left.left.label == 1

												with self.assertRaises(snake_case__    ):  # noqa: B017
																	t.put(1    )






							def 							__lowercase  (					self						:			Tuple    ):






												"""simple docstring"""



												__a:						Tuple		    =					self._get_binary_search_tree()

												__a:						Tuple		    =					t.search(6    )
												assert node.label == 6

												__a:						Tuple		    =					t.search(1_3    )
												assert node.label == 1_3

												with self.assertRaises(snake_case__    ):  # noqa: B017
																	t.search(2    )






							def 							__lowercase  (					self						:			List[str]    ):






												"""simple docstring"""



												__a:						List[Any]		    =					self._get_binary_search_tree()

												t.remove(1_3    )
												assert t.root is not None
												assert t.root.right is not None
												assert t.root.right.right is not None
												assert t.root.right.right.right is None
												assert t.root.right.right.left is None

												t.remove(7    )
												assert t.root.left is not None
												assert t.root.left.right is not None
												assert t.root.left.right.left is not None
												assert t.root.left.right.right is None
												assert t.root.left.right.left.label == 4

												t.remove(6    )
												assert t.root.left.left is not None
												assert t.root.left.right.right is not None
												assert t.root.left.left.label == 1
												assert t.root.left.right.label == 4
												assert t.root.left.right.right.label == 5
												assert t.root.left.right.left is None
												assert t.root.left.left.parent == t.root.left
												assert t.root.left.right.parent == t.root.left

												t.remove(3    )
												assert t.root is not None
												assert t.root.left.label == 4
												assert t.root.left.right.label == 5
												assert t.root.left.left.label == 1
												assert t.root.left.parent == t.root
												assert t.root.left.left.parent == t.root.left
												assert t.root.left.right.parent == t.root.left

												t.remove(4    )
												assert t.root.left is not None
												assert t.root.left.left is not None
												assert t.root.left.label == 5
												assert t.root.left.right is None
												assert t.root.left.left.label == 1
												assert t.root.left.parent == t.root
												assert t.root.left.left.parent == t.root.left






							def 							__lowercase  (					self						:			Optional[Any]    ):






												"""simple docstring"""



												__a:						Optional[int]		    =					self._get_binary_search_tree()

												t.remove(3    )
												assert t.root is not None
												assert t.root.left is not None
												assert t.root.left.left is not None
												assert t.root.left.right is not None
												assert t.root.left.right.left is not None
												assert t.root.left.right.right is not None
												assert t.root.left.label == 4
												assert t.root.left.right.label == 6
												assert t.root.left.left.label == 1
												assert t.root.left.right.right.label == 7
												assert t.root.left.right.left.label == 5
												assert t.root.left.parent == t.root
												assert t.root.left.right.parent == t.root.left
												assert t.root.left.left.parent == t.root.left
												assert t.root.left.right.left.parent == t.root.left.right






							def 							__lowercase  (					self						:			Optional[Any]    ):






												"""simple docstring"""



												__a:						Any		    =					self._get_binary_search_tree()
												t.empty()
												assert t.root is None






							def 							__lowercase  (					self						:			int    ):






												"""simple docstring"""



												__a:						List[str]		    =					self._get_binary_search_tree()
												assert not t.is_empty()

												t.empty()
												assert t.is_empty()






							def 							__lowercase  (					self						:			Optional[int]    ):






												"""simple docstring"""



												__a:						str		    =					self._get_binary_search_tree()

												assert t.exists(6    )
												assert not t.exists(-1    )






							def 							__lowercase  (					self						:			Optional[Any]    ):






												"""simple docstring"""



												__a:						Optional[Any]		    =					self._get_binary_search_tree()

												assert t.get_max_label() == 1_4

												t.empty()
												with self.assertRaises(snake_case__    ):  # noqa: B017
																	t.get_max_label()






							def 							__lowercase  (					self						:			str    ):






												"""simple docstring"""



												__a:						List[str]		    =					self._get_binary_search_tree()

												assert t.get_min_label() == 1

												t.empty()
												with self.assertRaises(snake_case__    ):  # noqa: B017
																	t.get_min_label()






							def 							__lowercase  (					self						:			Any    ):






												"""simple docstring"""



												__a:						Union[str, Any]		    =					self._get_binary_search_tree()

												__a:						Tuple		    =					[i.label for i in t.inorder_traversal()]
												assert inorder_traversal_nodes == [1, 3, 4, 5, 6, 7, 8, 1_0, 1_3, 1_4]






							def 							__lowercase  (					self						:			Optional[Any]    ):






												"""simple docstring"""



												__a:						List[str]		    =					self._get_binary_search_tree()

												__a:						Optional[Any]		    =					[i.label for i in t.preorder_traversal()]
												assert preorder_traversal_nodes == [8, 3, 1, 6, 4, 5, 7, 1_0, 1_4, 1_3]






def     lowerCamelCase_    (	)     ->			None:


					'''simple docstring'''
					__a:						List[Any]		    =					BinarySearchTree()
					t.put(8 )
					t.put(3 )
					t.put(6 )
					t.put(1 )
					t.put(10 )
					t.put(14 )
					t.put(13 )
					t.put(4 )
					t.put(7 )
					t.put(5 )

					print(
					    """
            8
           / \\
          3   10
         / \\    \\
        1   6    14
           / \\   /
          4   7 13
           \\
            5
        """ )

					print("""Label 6 exists:"""	,				t.exists(6 ) )
					print("""Label 13 exists:"""	,				t.exists(13 ) )
					print("""Label -1 exists:"""	,				t.exists(-1 ) )
					print("""Label 12 exists:"""	,				t.exists(12 ) )

					# Prints all the elements of the list in inorder traversal
					__a:						Union[str, Any]		    =					[i.label for i in t.inorder_traversal()]
					print("""Inorder traversal:"""	,				lowerCamelCase_ )

					# Prints all the elements of the list in preorder traversal
					__a:						Any		    =					[i.label for i in t.preorder_traversal()]
					print("""Preorder traversal:"""	,				lowerCamelCase_ )

					print("""Max. label:"""	,				t.get_max_label() )
					print("""Min. label:"""	,				t.get_min_label() )

					# Delete elements
					print("""\nDeleting elements 13, 10, 8, 3, 6, 14""" )
					print(
					    """
          4
         / \\
        1   7
             \\
              5
        """ )
					t.remove(13 )
					t.remove(10 )
					t.remove(8 )
					t.remove(3 )
					t.remove(6 )
					t.remove(14 )

					# Prints all the elements of the list in inorder traversal after delete
					__a:						Union[str, Any]		    =					[i.label for i in t.inorder_traversal()]
					print("""Inorder traversal after delete:"""	,				lowerCamelCase_ )

					# Prints all the elements of the list in preorder traversal after delete
					__a:						List[str]		    =					[i.label for i in t.preorder_traversal()]
					print("""Preorder traversal after delete:"""	,				lowerCamelCase_ )

					print("""Max. label:"""	,				t.get_max_label() )
					print("""Min. label:"""	,				t.get_min_label() )


if __name__ == "__main__":
			binary_search_tree_example()


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

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


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


					'''simple docstring'''
					if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
										__a:						int		    =					unittest.skip("""test requires Datasets, Faiss, psutil""" )(lowerCamelCase_ )
					return test_case




@require_distributed_retrieval
class    _snake_case  (						_lowercase       ):






							def 							__lowercase  (					self						:			Any    ):






												"""simple docstring"""



												__a:						Optional[Any]		    =					tempfile.mkdtemp()
												__a:						int		    =					8

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

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






							def 							__lowercase  (					self						:			Any    ):






												"""simple docstring"""



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






							def 							__lowercase  (					self						:			Any    ):






												"""simple docstring"""



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






							def 							__lowercase  (					self						:			Tuple    ):






												"""simple docstring"""



												shutil.rmtree(self.tmpdirname    )






							def 							__lowercase  (					self						:			int    ):






												"""simple docstring"""



												__a:						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 							__lowercase  (					self						:			Any						, snake_case__						:			bool						, snake_case__						:			Dict=1_2_3_4_5    ):






												"""simple docstring"""



												__a:						str		    =					self.get_dummy_dataset()
												__a:						int		    =					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:
																	__a:						List[str]		    =					dataset
																	__a:						List[Any]		    =					RagPyTorchDistributedRetriever(
																	    snake_case__						, question_encoder_tokenizer=self.get_dpr_tokenizer()						, generator_tokenizer=self.get_bart_tokenizer()						, )
																	if init_retrieval:
																						retriever.init_retrieval(snake_case__    )
												return retriever






							def 							__lowercase  (					self						:			Dict						, snake_case__						:			bool    ):






												"""simple docstring"""



												ray.init(local_mode=snake_case__    )
												__a:						int		    =					RagConfig(
												    retrieval_vector_size=self.retrieval_vector_size						, question_encoder=DPRConfig().to_dict()						, generator=BartConfig().to_dict()						, )
												__a:						List[str]		    =					ray.remote(snake_case__    )
												__a:						Tuple		    =					[remote_cls.remote() for _ in range(1    )]
												with patch("""transformers.models.rag.retrieval_rag.load_dataset"""    ) as mock_load_dataset:
																	__a:						int		    =					self.get_dummy_dataset()
																	__a:						List[Any]		    =					RagRayDistributedRetriever(
																	    snake_case__						, question_encoder_tokenizer=self.get_dpr_tokenizer()						, generator_tokenizer=self.get_bart_tokenizer()						, retrieval_workers=snake_case__						, )
																	if init_retrieval:
																						retriever.init_retrieval()
												return retriever






							def 							__lowercase  (					self						:			str						, snake_case__						:			bool						, snake_case__						:			bool						, snake_case__						:			Union[str, Any]=1_2_3_4_5    ):






												"""simple docstring"""



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






							def 							__lowercase  (					self						:			str						, snake_case__						:			bool						, snake_case__						:			bool    ):






												"""simple docstring"""



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






							def 							__lowercase  (					self						:			str						, snake_case__						:			RagRetriever						, snake_case__						:			np.array						, snake_case__						:			int    ):






												"""simple docstring"""



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






												"""simple docstring"""



												__a:						int		    =					1
												__a:						Tuple		    =					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=snake_case__    )						, snake_case__						, snake_case__    )






							def 							__lowercase  (					self						:			str    ):






												"""simple docstring"""



												__a:						Tuple		    =					1
												__a:						Tuple		    =					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=snake_case__						, from_disk=snake_case__    )						, snake_case__						, snake_case__						, )






							def 							__lowercase  (					self						:			List[Any]    ):






												"""simple docstring"""



												__a:						List[Any]		    =					1
												__a:						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=snake_case__						, from_disk=snake_case__    )						, snake_case__						, snake_case__						, )






							@require_ray
							def 							__lowercase  (					self						:			str    ):






												"""simple docstring"""



												__a:						int		    =					1
												__a:						str		    =					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=snake_case__    )						, snake_case__						, snake_case__    )
												ray.shutdown()






							@require_ray
							def 							__lowercase  (					self						:			Dict    ):






												"""simple docstring"""



												__a:						List[Any]		    =					1
												__a:						List[str]		    =					np.array(
												    [np.ones(self.retrieval_vector_size    ), -np.ones(self.retrieval_vector_size    )]						, dtype=np.floataa    )
												with self.assertRaises(snake_case__    ):
																	self.distributed_retriever_check(
																	    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=snake_case__						, from_disk=snake_case__    )						, snake_case__						, snake_case__						, )
												ray.shutdown()






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






												"""simple docstring"""



												__a:						Tuple		    =					1
												__a:						str		    =					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=snake_case__						, from_disk=snake_case__    )						, snake_case__						, snake_case__    )
												ray.shutdown()
