





snake_case								=							{
    """Pillow""": """Pillow""",
    """accelerate""": """accelerate>=0.11.0""",
    """compel""": """compel==0.1.8""",
    """black""": """black~=23.1""",
    """datasets""": """datasets""",
    """filelock""": """filelock""",
    """flax""": """flax>=0.4.1""",
    """hf-doc-builder""": """hf-doc-builder>=0.3.0""",
    """huggingface-hub""": """huggingface-hub>=0.13.2""",
    """requests-mock""": """requests-mock==1.10.0""",
    """importlib_metadata""": """importlib_metadata""",
    """invisible-watermark""": """invisible-watermark""",
    """isort""": """isort>=5.5.4""",
    """jax""": """jax>=0.2.8,!=0.3.2""",
    """jaxlib""": """jaxlib>=0.1.65""",
    """Jinja2""": """Jinja2""",
    """k-diffusion""": """k-diffusion>=0.0.12""",
    """torchsde""": """torchsde""",
    """note_seq""": """note_seq""",
    """librosa""": """librosa""",
    """numpy""": """numpy""",
    """omegaconf""": """omegaconf""",
    """parameterized""": """parameterized""",
    """protobuf""": """protobuf>=3.20.3,<4""",
    """pytest""": """pytest""",
    """pytest-timeout""": """pytest-timeout""",
    """pytest-xdist""": """pytest-xdist""",
    """ruff""": """ruff>=0.0.241""",
    """safetensors""": """safetensors""",
    """sentencepiece""": """sentencepiece>=0.1.91,!=0.1.92""",
    """scipy""": """scipy""",
    """onnx""": """onnx""",
    """regex""": """regex!=2019.12.17""",
    """requests""": """requests""",
    """tensorboard""": """tensorboard""",
    """torch""": """torch>=1.4""",
    """torchvision""": """torchvision""",
    """transformers""": """transformers>=4.25.1""",
    """urllib3""": """urllib3<=2.0.0""",
}


"""simple docstring"""




import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pytest

from datasets.arrow_dataset import Dataset
from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex

from .utils import require_elasticsearch, require_faiss


A_				:      List[str]           =			pytest.mark.integration





@require_faiss
class  lowerCamelCase    (A__      ):

			def        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     Tuple:
						SCREAMING_SNAKE_CASE__										=   Dataset.from_dict({"""filename""": ["""my_name-train""" + """_""" + str(__UpperCAmelCase							) for x in np.arange(3_0							).tolist()]}							)
						return dset

			def        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     Tuple:
						import faiss

						SCREAMING_SNAKE_CASE__										=   self._create_dummy_dataset()
						SCREAMING_SNAKE_CASE__										=   dset.map(
						    lambda __UpperCAmelCase     ,					__UpperCAmelCase							: {"vecs": i * np.ones(5     ,					dtype=np.floataa							)}     ,					with_indices=__UpperCAmelCase     ,					keep_in_memory=__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   dset.add_faiss_index("""vecs"""     ,					batch_size=1_0_0     ,					metric_type=faiss.METRIC_INNER_PRODUCT							)
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   dset.get_nearest_examples("""vecs"""     ,					np.ones(5     ,					dtype=np.floataa							)							)
						self.assertEqual(examples["""filename"""][0]     ,					"""my_name-train_29"""							)
						dset.drop_index("""vecs"""							)

			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     str:
						import faiss

						SCREAMING_SNAKE_CASE__										=   self._create_dummy_dataset()
						dset.add_faiss_index_from_external_arrays(
						    external_arrays=np.ones((3_0, 5)							) * np.arange(3_0							).reshape(-1     ,					1							)     ,					index_name="""vecs"""     ,					batch_size=1_0_0     ,					metric_type=faiss.METRIC_INNER_PRODUCT     ,					)
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   dset.get_nearest_examples("""vecs"""     ,					np.ones(5     ,					dtype=np.floataa							)							)
						self.assertEqual(examples["""filename"""][0]     ,					"""my_name-train_29"""							)

			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     Any:
						import faiss

						SCREAMING_SNAKE_CASE__										=   self._create_dummy_dataset()
						dset.add_faiss_index_from_external_arrays(
						    external_arrays=np.ones((3_0, 5)							) * np.arange(3_0							).reshape(-1     ,					1							)     ,					index_name="""vecs"""     ,					metric_type=faiss.METRIC_INNER_PRODUCT     ,					)

						# Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
						# ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
						# see https://bugs.python.org/issue14243 and
						# https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
						with tempfile.NamedTemporaryFile(delete=__UpperCAmelCase							) as tmp_file:
									dset.save_faiss_index("""vecs"""     ,					tmp_file.name							)
									dset.load_faiss_index("""vecs2"""     ,					tmp_file.name							)
						os.unlink(tmp_file.name							)

						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   dset.get_nearest_examples("""vecs2"""     ,					np.ones(5     ,					dtype=np.floataa							)							)
						self.assertEqual(examples["""filename"""][0]     ,					"""my_name-train_29"""							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[int]							) ->     int:
						SCREAMING_SNAKE_CASE__										=   self._create_dummy_dataset()
						dset.add_faiss_index_from_external_arrays(
						    external_arrays=np.ones((3_0, 5)							) * np.arange(3_0							).reshape(-1     ,					1							)     ,					index_name="""vecs"""							)
						dset.drop_index("""vecs"""							)
						self.assertRaises(__UpperCAmelCase     ,					partial(dset.get_nearest_examples     ,					"""vecs2"""     ,					np.ones(5     ,					dtype=np.floataa							)							)							)





			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     str:
						from elasticsearch import Elasticsearch

						SCREAMING_SNAKE_CASE__										=   self._create_dummy_dataset()
						with patch("""elasticsearch.Elasticsearch.search"""							) as mocked_search, patch(
						    """elasticsearch.client.IndicesClient.create"""							) as mocked_index_create, patch("""elasticsearch.helpers.streaming_bulk"""							) as mocked_bulk:
									SCREAMING_SNAKE_CASE__										=   {"""acknowledged""": True}
									mocked_bulk.return_value([(True, None)] * 3_0							)
									SCREAMING_SNAKE_CASE__										=   {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 2_9}]}}
									SCREAMING_SNAKE_CASE__										=   Elasticsearch()

									dset.add_elasticsearch_index("""filename"""     ,					es_client=__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   dset.get_nearest_examples("""filename"""     ,					"""my_name-train_29"""							)
									self.assertEqual(examples["""filename"""][0]     ,					"""my_name-train_29"""							)





@require_faiss
class  lowerCamelCase    (A__      ):

			def        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     List[str]:
						import faiss

						SCREAMING_SNAKE_CASE__										=   FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT							)

						# add vectors
						index.add_vectors(np.eye(5     ,					dtype=np.floataa							)							)
						self.assertIsNotNone(index.faiss_index							)
						self.assertEqual(index.faiss_index.ntotal     ,					5							)
						index.add_vectors(np.zeros((5, 5)     ,					dtype=np.floataa							)							)
						self.assertEqual(index.faiss_index.ntotal     ,					1_0							)

						# single query
						SCREAMING_SNAKE_CASE__										=   np.zeros(5     ,					dtype=np.floataa							)
						SCREAMING_SNAKE_CASE__										=   1
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   index.search(__UpperCAmelCase							)
						self.assertRaises(__UpperCAmelCase     ,					index.search     ,					query.reshape(-1     ,					1							)							)
						self.assertGreater(scores[0]     ,					0							)
						self.assertEqual(indices[0]     ,					1							)

						# batched queries
						SCREAMING_SNAKE_CASE__										=   np.eye(5     ,					dtype=np.floataa							)[::-1]
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   index.search_batch(__UpperCAmelCase							)
						self.assertRaises(__UpperCAmelCase     ,					index.search_batch     ,					queries[0]							)
						SCREAMING_SNAKE_CASE__										=   [scores[0] for scores in total_scores]
						SCREAMING_SNAKE_CASE__										=   [indices[0] for indices in total_indices]
						self.assertGreater(np.min(__UpperCAmelCase							)     ,					0							)
						self.assertListEqual([4, 3, 2, 1, 0]     ,					__UpperCAmelCase							)

			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     str:
						import faiss

						SCREAMING_SNAKE_CASE__										=   FaissIndex(string_factory="""Flat"""							)
						index.add_vectors(np.eye(5     ,					dtype=np.floataa							)							)
						self.assertIsInstance(index.faiss_index     ,					faiss.IndexFlat							)
						SCREAMING_SNAKE_CASE__										=   FaissIndex(string_factory="""LSH"""							)
						index.add_vectors(np.eye(5     ,					dtype=np.floataa							)							)
						self.assertIsInstance(index.faiss_index     ,					faiss.IndexLSH							)
						with self.assertRaises(__UpperCAmelCase							):
									SCREAMING_SNAKE_CASE__										=   FaissIndex(string_factory="""Flat"""     ,					custom_index=faiss.IndexFlat(5							)							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     Any:
						import faiss

						SCREAMING_SNAKE_CASE__										=   faiss.IndexFlat(5							)
						SCREAMING_SNAKE_CASE__										=   FaissIndex(custom_index=__UpperCAmelCase							)
						index.add_vectors(np.eye(5     ,					dtype=np.floataa							)							)
						self.assertIsInstance(index.faiss_index     ,					faiss.IndexFlat							)





			def        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]							) ->     str:
						import faiss

						SCREAMING_SNAKE_CASE__										=   FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT							)
						index.add_vectors(np.eye(5     ,					dtype=np.floataa							)							)

						# Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
						# ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
						# see https://bugs.python.org/issue14243 and
						# https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
						with tempfile.NamedTemporaryFile(delete=__UpperCAmelCase							) as tmp_file:
									index.save(tmp_file.name							)
									SCREAMING_SNAKE_CASE__										=   FaissIndex.load(tmp_file.name							)
						os.unlink(tmp_file.name							)

						SCREAMING_SNAKE_CASE__										=   np.zeros(5     ,					dtype=np.floataa							)
						SCREAMING_SNAKE_CASE__										=   1
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   index.search(__UpperCAmelCase							)
						self.assertGreater(scores[0]     ,					0							)
						self.assertEqual(indices[0]     ,					1							)

@require_faiss
def 						A							(  snake_case__     ):




			'''simple docstring'''


			import faiss

			SCREAMING_SNAKE_CASE__										=   FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT     )
			index.add_vectors(np.eye(5		,					dtype=np.floataa     )     )

			SCREAMING_SNAKE_CASE__										=   """index.faiss"""
			SCREAMING_SNAKE_CASE__										=   f"""mock://{index_name}"""
			index.save(snake_case__		,					storage_options=mockfs.storage_options     )
			SCREAMING_SNAKE_CASE__										=   FaissIndex.load(snake_case__		,					storage_options=mockfs.storage_options     )

			SCREAMING_SNAKE_CASE__										=   np.zeros(5		,					dtype=np.floataa     )
			SCREAMING_SNAKE_CASE__										=   1
			SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   index.search(snake_case__     )
			assert scores[0] > 0
			assert indices[0] == 1





@require_elasticsearch
class  lowerCamelCase    (A__      ):

			def        SCREAMING_SNAKE_CASE					(       self					:    Tuple							) ->     Dict:
						from elasticsearch import Elasticsearch

						with patch("""elasticsearch.Elasticsearch.search"""							) as mocked_search, patch(
						    """elasticsearch.client.IndicesClient.create"""							) as mocked_index_create, patch("""elasticsearch.helpers.streaming_bulk"""							) as mocked_bulk:
									SCREAMING_SNAKE_CASE__										=   Elasticsearch()
									SCREAMING_SNAKE_CASE__										=   {"""acknowledged""": True}
									SCREAMING_SNAKE_CASE__										=   ElasticSearchIndex(es_client=__UpperCAmelCase							)
									mocked_bulk.return_value([(True, None)] * 3							)
									index.add_documents(["""foo""", """bar""", """foobar"""]							)

									# single query
									SCREAMING_SNAKE_CASE__										=   """foo"""
									SCREAMING_SNAKE_CASE__										=   {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 0}]}}
									SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   index.search(__UpperCAmelCase							)
									self.assertEqual(scores[0]     ,					1							)
									self.assertEqual(indices[0]     ,					0							)

									# single query with timeout
									SCREAMING_SNAKE_CASE__										=   """foo"""
									SCREAMING_SNAKE_CASE__										=   {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 0}]}}
									SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   index.search(__UpperCAmelCase     ,					request_timeout=3_0							)
									self.assertEqual(scores[0]     ,					1							)
									self.assertEqual(indices[0]     ,					0							)

									# batched queries
									SCREAMING_SNAKE_CASE__										=   ["""foo""", """bar""", """foobar"""]
									SCREAMING_SNAKE_CASE__										=   {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 1}]}}
									SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   index.search_batch(__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__										=   [scores[0] for scores in total_scores]
									SCREAMING_SNAKE_CASE__										=   [indices[0] for indices in total_indices]
									self.assertGreater(np.min(__UpperCAmelCase							)     ,					0							)
									self.assertListEqual([1, 1, 1]     ,					__UpperCAmelCase							)

									# batched queries with timeout
									SCREAMING_SNAKE_CASE__										=   ["""foo""", """bar""", """foobar"""]
									SCREAMING_SNAKE_CASE__										=   {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 1}]}}
									SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   index.search_batch(__UpperCAmelCase     ,					request_timeout=3_0							)
									SCREAMING_SNAKE_CASE__										=   [scores[0] for scores in total_scores]
									SCREAMING_SNAKE_CASE__										=   [indices[0] for indices in total_indices]
									self.assertGreater(np.min(__UpperCAmelCase							)     ,					0							)
									self.assertListEqual([1, 1, 1]     ,					__UpperCAmelCase							)


