






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


UpperCamelCase     :   Optional[int]	      =		pytest.mark.integration




@require_faiss
class 		UpperCamelCase__     (_snake_case      ):




		'''simple docstring'''





		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	Dataset.from_dict({"""filename""": ["""my_name-train""" + """_""" + str(__snake_case							) for x in np.arange(30							).tolist()]}							)
								return dset






		def       UpperCamelCase_       (   self							):
								import faiss

								lowerCamelCase__								=	self._create_dummy_dataset()
								lowerCamelCase__								=	dset.map(
								    lambda _lowerCAmelCase ,_lowerCAmelCase							: {"vecs": i * np.ones(5 ,dtype=np.floataa							)} ,with_indices=__snake_case ,keep_in_memory=__snake_case							)
								lowerCamelCase__								=	dset.add_faiss_index("""vecs""" ,batch_size=1_00 ,metric_type=faiss.METRIC_INNER_PRODUCT							)
								lowerCamelCase__								=	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       UpperCamelCase_       (   self							):
								import faiss

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






		def       UpperCamelCase_       (   self							):
								import faiss

								lowerCamelCase__								=	self._create_dummy_dataset()
								dset.add_faiss_index_from_external_arrays(
								    external_arrays=np.ones((30, 5)							) * np.arange(30							).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=__snake_case							) as tmp_file:
														dset.save_faiss_index("""vecs""" ,tmp_file.name							)
														dset.load_faiss_index("""vecs2""" ,tmp_file.name							)
								os.unlink(tmp_file.name							)

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






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self._create_dummy_dataset()
								dset.add_faiss_index_from_external_arrays(
								    external_arrays=np.ones((30, 5)							) * np.arange(30							).reshape(-1 ,1							) ,index_name="""vecs"""							)
								dset.drop_index("""vecs"""							)
								self.assertRaises(__snake_case ,partial(dset.get_nearest_examples ,"""vecs2""" ,np.ones(5 ,dtype=np.floataa							)							)							)






		def       UpperCamelCase_       (   self							):
								from elasticsearch import Elasticsearch

								lowerCamelCase__								=	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:
														lowerCamelCase__								=	{"""acknowledged""": True}
														mocked_bulk.return_value([(True, None)] * 30							)
														lowerCamelCase__								=	{"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 29}]}}
														lowerCamelCase__								=	Elasticsearch()

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




@require_faiss
class 		UpperCamelCase__     (_snake_case      ):




		'''simple docstring'''





		def       UpperCamelCase_       (   self							):
								import faiss

								lowerCamelCase__								=	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 ,10							)

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

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






		def       UpperCamelCase_       (   self							):
								import faiss

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






		def       UpperCamelCase_       (   self							):
								import faiss

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






		def       UpperCamelCase_       (   self							):
								import faiss

								lowerCamelCase__								=	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=__snake_case							) as tmp_file:
														index.save(tmp_file.name							)
														lowerCamelCase__								=	FaissIndex.load(tmp_file.name							)
								os.unlink(tmp_file.name							)

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



@require_faiss
def     A__      (					__lowerCAmelCase		:      Any    ):
						import faiss

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

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

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





@require_elasticsearch
class 		UpperCamelCase__     (_snake_case      ):




		'''simple docstring'''





		def       UpperCamelCase_       (   self							):
								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:
														lowerCamelCase__								=	Elasticsearch()
														lowerCamelCase__								=	{"""acknowledged""": True}
														lowerCamelCase__								=	ElasticSearchIndex(es_client=__snake_case							)
														mocked_bulk.return_value([(True, None)] * 3							)
														index.add_documents(["""foo""", """bar""", """foobar"""]							)

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

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

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

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





'''simple docstring'''


from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


UpperCAmelCase_						:  int				  =						logging.get_logger(__name__)


if is_vision_available():
				import PIL





class        lowercase__    (		_snake_case   ):





							'''simple docstring'''


							A_				: int       =    ["""pixel_values"""]






							def __init__(       self       ,   __snake_case = True       ,   __snake_case = None       ,   __snake_case = PILImageResampling.BICUBIC       ,   __snake_case = True       ,   __snake_case = None       ,   __snake_case = True       ,   __snake_case = 1 / 255       ,   __snake_case = True       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = True       ,   **__snake_case       ,   ):
											super().__init__(**__snake_case    )
											_SCREAMING_SNAKE_CASE					:			int							= size if size is not None else {"""shortest_edge""": 224}
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= get_size_dict(__snake_case       ,   default_to_square=__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
											_SCREAMING_SNAKE_CASE					:			Any							= get_size_dict(__snake_case       ,   default_to_square=__snake_case       ,   param_name="""crop_size"""    )

											_SCREAMING_SNAKE_CASE					:			str							= do_resize
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= size
											_SCREAMING_SNAKE_CASE					:			List[str]							= resample
											_SCREAMING_SNAKE_CASE					:			str							= do_center_crop
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= crop_size
											_SCREAMING_SNAKE_CASE					:			Tuple							= do_rescale
											_SCREAMING_SNAKE_CASE					:			Dict							= rescale_factor
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= do_normalize
											_SCREAMING_SNAKE_CASE					:			List[Any]							= image_mean if image_mean is not None else OPENAI_CLIP_MEAN
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= image_std if image_std is not None else OPENAI_CLIP_STD
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= do_convert_rgb






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case = PILImageResampling.BICUBIC       ,   __snake_case = None       ,   **__snake_case       ,   ):
											_SCREAMING_SNAKE_CASE					:			int							= get_size_dict(__snake_case       ,   default_to_square=__snake_case    )
											if "shortest_edge" not in size:
															raise ValueError(f"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""    )
											_SCREAMING_SNAKE_CASE					:			Any							= get_resize_output_image_size(__snake_case       ,   size=size["""shortest_edge"""]       ,   default_to_square=__snake_case    )
											return resize(__snake_case       ,   size=__snake_case       ,   resample=__snake_case       ,   data_format=__snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case = None       ,   **__snake_case       ,   ):
											_SCREAMING_SNAKE_CASE					:			List[str]							= get_size_dict(__snake_case    )
											if "height" not in size or "width" not in size:
															raise ValueError(f"""The `size` parameter must contain the keys (height, width). Got {size.keys()}"""    )
											return center_crop(__snake_case       ,   size=(size["""height"""], size["""width"""])       ,   data_format=__snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case = None       ,   **__snake_case       ,   ):
											return rescale(__snake_case       ,   scale=__snake_case       ,   data_format=__snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case = None       ,   **__snake_case       ,   ):
											return normalize(__snake_case       ,   mean=__snake_case       ,   std=__snake_case       ,   data_format=__snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = None       ,   __snake_case = ChannelDimension.FIRST       ,   **__snake_case       ,   ):
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= do_resize if do_resize is not None else self.do_resize
											_SCREAMING_SNAKE_CASE					:			Any							= size if size is not None else self.size
											_SCREAMING_SNAKE_CASE					:			List[str]							= get_size_dict(__snake_case       ,   param_name="""size"""       ,   default_to_square=__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Any							= resample if resample is not None else self.resample
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= do_center_crop if do_center_crop is not None else self.do_center_crop
											_SCREAMING_SNAKE_CASE					:			str							= crop_size if crop_size is not None else self.crop_size
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= get_size_dict(__snake_case       ,   param_name="""crop_size"""       ,   default_to_square=__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= do_rescale if do_rescale is not None else self.do_rescale
											_SCREAMING_SNAKE_CASE					:			int							= rescale_factor if rescale_factor is not None else self.rescale_factor
											_SCREAMING_SNAKE_CASE					:			Tuple							= do_normalize if do_normalize is not None else self.do_normalize
											_SCREAMING_SNAKE_CASE					:			List[str]							= image_mean if image_mean is not None else self.image_mean
											_SCREAMING_SNAKE_CASE					:			List[Any]							= image_std if image_std is not None else self.image_std
											_SCREAMING_SNAKE_CASE					:			Any							= do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= make_list_of_images(__snake_case    )

											if not valid_images(__snake_case    ):
															raise ValueError(
															    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
															    """torch.Tensor, tf.Tensor or jax.ndarray."""    )

											if do_resize and size is None:
															raise ValueError("""Size must be specified if do_resize is True."""    )

											if do_center_crop and crop_size is None:
															raise ValueError("""Crop size must be specified if do_center_crop is True."""    )

											if do_rescale and rescale_factor is None:
															raise ValueError("""Rescale factor must be specified if do_rescale is True."""    )

											if do_normalize and (image_mean is None or image_std is None):
															raise ValueError("""Image mean and std must be specified if do_normalize is True."""    )

											# PIL RGBA images are converted to RGB
											if do_convert_rgb:
															_SCREAMING_SNAKE_CASE					:			Optional[int]							= [convert_to_rgb(__snake_case    ) for image in images]

											# All transformations expect numpy arrays.
											_SCREAMING_SNAKE_CASE					:			str							= [to_numpy_array(__snake_case    ) for image in images]

											if do_resize:
															_SCREAMING_SNAKE_CASE					:			Optional[Any]							= [self.resize(image=__snake_case       ,   size=__snake_case       ,   resample=__snake_case    ) for image in images]

											if do_center_crop:
															_SCREAMING_SNAKE_CASE					:			Tuple							= [self.center_crop(image=__snake_case       ,   size=__snake_case    ) for image in images]

											if do_rescale:
															_SCREAMING_SNAKE_CASE					:			List[Any]							= [self.rescale(image=__snake_case       ,   scale=__snake_case    ) for image in images]

											if do_normalize:
															_SCREAMING_SNAKE_CASE					:			Dict							= [self.normalize(image=__snake_case       ,   mean=__snake_case       ,   std=__snake_case    ) for image in images]

											_SCREAMING_SNAKE_CASE					:			str							= [to_channel_dimension_format(__snake_case       ,   __snake_case    ) for image in images]

											_SCREAMING_SNAKE_CASE					:			int							= {"""pixel_values""": images}
											return BatchFeature(data=__snake_case       ,   tensor_type=__snake_case    )




