





from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block




@dataclass
class 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       42




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Any				,					__A     :str=3				,					__A     :List[Any]=3				,					__A     :str=("DownEncoderBlock2D",)				,					__A     :Optional[int]=(64,)				,					__A     :str=2				,					__A     :Optional[Any]=32				,					__A     :int="silu"				,					__A     :List[Any]=True				,					) -> Optional[Any]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   layers_per_block

									SCREAMING_SNAKE_CASE__								=   torch.nn.Convad(
									    __A				,					block_out_channels[0]				,					kernel_size=3				,					stride=1				,					padding=1				,					)

									SCREAMING_SNAKE_CASE__								=   None
									SCREAMING_SNAKE_CASE__								=   nn.ModuleList([]		)

									# down
									SCREAMING_SNAKE_CASE__								=   block_out_channels[0]
									for i, down_block_type in enumerate(__A		):
											SCREAMING_SNAKE_CASE__								=   output_channel
											SCREAMING_SNAKE_CASE__								=   block_out_channels[i]
											SCREAMING_SNAKE_CASE__								=   i == len(__A		) - 1

											SCREAMING_SNAKE_CASE__								=   get_down_block(
											    __A				,					num_layers=self.layers_per_block				,					in_channels=__A				,					out_channels=__A				,					add_downsample=not is_final_block				,					resnet_eps=1E-6				,					downsample_padding=0				,					resnet_act_fn=__A				,					resnet_groups=__A				,					attention_head_dim=__A				,					temb_channels=__A				,					)
											self.down_blocks.append(__A		)

									# mid
									SCREAMING_SNAKE_CASE__								=   UNetMidBlockaD(
									    in_channels=block_out_channels[-1]				,					resnet_eps=1E-6				,					resnet_act_fn=__A				,					output_scale_factor=1				,					resnet_time_scale_shift="""default"""				,					attention_head_dim=block_out_channels[-1]				,					resnet_groups=__A				,					temb_channels=__A				,					)

									# out
									SCREAMING_SNAKE_CASE__								=   nn.GroupNorm(num_channels=block_out_channels[-1]				,					num_groups=__A				,					eps=1E-6		)
									SCREAMING_SNAKE_CASE__								=   nn.SiLU()

									SCREAMING_SNAKE_CASE__								=   2 * out_channels if double_z else out_channels
									SCREAMING_SNAKE_CASE__								=   nn.Convad(block_out_channels[-1]				,					__A				,					3				,					padding=1		)

									SCREAMING_SNAKE_CASE__								=   False


							def   _snake_case	(				self     :Optional[int]				,					__A     :Tuple		) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   x
									SCREAMING_SNAKE_CASE__								=   self.conv_in(__A		)

									if self.training and self.gradient_checkpointing:

											def create_custom_forward(__A     :Tuple		):
													def custom_forward(*__A     :List[Any]		):
															return module(*__A		)

													return custom_forward

											# down
											if is_torch_version(""">="""				,					"""1.11.0"""		):
													for down_block in self.down_blocks:
															SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(
															    create_custom_forward(__A		)				,					__A				,					use_reentrant=__A		)
													# middle
													SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(
													    create_custom_forward(self.mid_block		)				,					__A				,					use_reentrant=__A		)
											else:
													for down_block in self.down_blocks:
															SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(create_custom_forward(__A		)				,					__A		)
													# middle
													SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block		)				,					__A		)

									else:
											# down
											for down_block in self.down_blocks:
													SCREAMING_SNAKE_CASE__								=   down_block(__A		)

											# middle
											SCREAMING_SNAKE_CASE__								=   self.mid_block(__A		)

									# post-process
									SCREAMING_SNAKE_CASE__								=   self.conv_norm_out(__A		)
									SCREAMING_SNAKE_CASE__								=   self.conv_act(__A		)
									SCREAMING_SNAKE_CASE__								=   self.conv_out(__A		)

									return sample




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Any				,					__A     :List[Any]=3				,					__A     :str=3				,					__A     :Tuple=("UpDecoderBlock2D",)				,					__A     :Tuple=(64,)				,					__A     :Any=2				,					__A     :int=32				,					__A     :Optional[int]="silu"				,					__A     :Optional[Any]="group"				,					) -> Optional[int]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   layers_per_block

									SCREAMING_SNAKE_CASE__								=   nn.Convad(
									    __A				,					block_out_channels[-1]				,					kernel_size=3				,					stride=1				,					padding=1				,					)

									SCREAMING_SNAKE_CASE__								=   None
									SCREAMING_SNAKE_CASE__								=   nn.ModuleList([]		)

									SCREAMING_SNAKE_CASE__								=   in_channels if norm_type == """spatial""" else None

									# mid
									SCREAMING_SNAKE_CASE__								=   UNetMidBlockaD(
									    in_channels=block_out_channels[-1]				,					resnet_eps=1E-6				,					resnet_act_fn=__A				,					output_scale_factor=1				,					resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type				,					attention_head_dim=block_out_channels[-1]				,					resnet_groups=__A				,					temb_channels=__A				,					)

									# up
									SCREAMING_SNAKE_CASE__								=   list(reversed(__A		)		)
									SCREAMING_SNAKE_CASE__								=   reversed_block_out_channels[0]
									for i, up_block_type in enumerate(__A		):
											SCREAMING_SNAKE_CASE__								=   output_channel
											SCREAMING_SNAKE_CASE__								=   reversed_block_out_channels[i]

											SCREAMING_SNAKE_CASE__								=   i == len(__A		) - 1

											SCREAMING_SNAKE_CASE__								=   get_up_block(
											    __A				,					num_layers=self.layers_per_block + 1				,					in_channels=__A				,					out_channels=__A				,					prev_output_channel=__A				,					add_upsample=not is_final_block				,					resnet_eps=1E-6				,					resnet_act_fn=__A				,					resnet_groups=__A				,					attention_head_dim=__A				,					temb_channels=__A				,					resnet_time_scale_shift=__A				,					)
											self.up_blocks.append(__A		)
											SCREAMING_SNAKE_CASE__								=   output_channel

									# out
									if norm_type == "spatial":
											SCREAMING_SNAKE_CASE__								=   SpatialNorm(block_out_channels[0]				,					__A		)
									else:
											SCREAMING_SNAKE_CASE__								=   nn.GroupNorm(num_channels=block_out_channels[0]				,					num_groups=__A				,					eps=1E-6		)
									SCREAMING_SNAKE_CASE__								=   nn.SiLU()
									SCREAMING_SNAKE_CASE__								=   nn.Convad(block_out_channels[0]				,					__A				,					3				,					padding=1		)

									SCREAMING_SNAKE_CASE__								=   False


							def   _snake_case	(				self     :str				,					__A     :Tuple				,					__A     :int=None		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   z
									SCREAMING_SNAKE_CASE__								=   self.conv_in(__A		)

									SCREAMING_SNAKE_CASE__								=   next(iter(self.up_blocks.parameters()		)		).dtype
									if self.training and self.gradient_checkpointing:

											def create_custom_forward(__A     :Tuple		):
													def custom_forward(*__A     :List[Any]		):
															return module(*__A		)

													return custom_forward

											if is_torch_version(""">="""				,					"""1.11.0"""		):
													# middle
													SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(
													    create_custom_forward(self.mid_block		)				,					__A				,					__A				,					use_reentrant=__A		)
													SCREAMING_SNAKE_CASE__								=   sample.to(__A		)

													# up
													for up_block in self.up_blocks:
															SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(
															    create_custom_forward(__A		)				,					__A				,					__A				,					use_reentrant=__A		)
											else:
													# middle
													SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(
													    create_custom_forward(self.mid_block		)				,					__A				,					__A		)
													SCREAMING_SNAKE_CASE__								=   sample.to(__A		)

													# up
													for up_block in self.up_blocks:
															SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(create_custom_forward(__A		)				,					__A				,					__A		)
									else:
											# middle
											SCREAMING_SNAKE_CASE__								=   self.mid_block(__A				,					__A		)
											SCREAMING_SNAKE_CASE__								=   sample.to(__A		)

											# up
											for up_block in self.up_blocks:
													SCREAMING_SNAKE_CASE__								=   up_block(__A				,					__A		)

        # post-process
									if latent_embeds is None:
											SCREAMING_SNAKE_CASE__								=   self.conv_norm_out(__A		)
									else:
											SCREAMING_SNAKE_CASE__								=   self.conv_norm_out(__A				,					__A		)
									SCREAMING_SNAKE_CASE__								=   self.conv_act(__A		)
									SCREAMING_SNAKE_CASE__								=   self.conv_out(__A		)

									return sample




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Optional[int]				,					__A     :Optional[Any]				,					__A     :List[Any]				,					__A     :str				,					__A     :List[str]=None				,					__A     :Optional[Any]="random"				,					__A     :Optional[int]=False				,					__A     :int=True		) -> Optional[int]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   n_e
									SCREAMING_SNAKE_CASE__								=   vq_embed_dim
									SCREAMING_SNAKE_CASE__								=   beta
									SCREAMING_SNAKE_CASE__								=   legacy

									SCREAMING_SNAKE_CASE__								=   nn.Embedding(self.n_e				,					self.vq_embed_dim		)
									self.embedding.weight.data.uniform_(-1.0 / self.n_e				,					1.0 / self.n_e		)

									SCREAMING_SNAKE_CASE__								=   remap
									if self.remap is not None:
											self.register_buffer("""used"""				,					torch.tensor(np.load(self.remap		)		)		)
											SCREAMING_SNAKE_CASE__								=   self.used.shape[0]
											SCREAMING_SNAKE_CASE__								=   unknown_index  # "random" or "extra" or integer
											if self.unknown_index == "extra":
													SCREAMING_SNAKE_CASE__								=   self.re_embed
													SCREAMING_SNAKE_CASE__								=   self.re_embed + 1
											print(
											    f'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
											    f'''Using {self.unknown_index} for unknown indices.'''		)
									else:
											SCREAMING_SNAKE_CASE__								=   n_e

									SCREAMING_SNAKE_CASE__								=   sane_index_shape


							def   _snake_case	(				self     :Tuple				,					__A     :List[Any]		) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   inds.shape
									assert len(__A		) > 1
									SCREAMING_SNAKE_CASE__								=   inds.reshape(ishape[0]				,					-1		)
									SCREAMING_SNAKE_CASE__								=   self.used.to(__A		)
									SCREAMING_SNAKE_CASE__								=   (inds[:, :, None] == used[None, None, ...]).long()
									SCREAMING_SNAKE_CASE__								=   match.argmax(-1		)
									SCREAMING_SNAKE_CASE__								=   match.sum(2		) < 1
									if self.unknown_index == "random":
											SCREAMING_SNAKE_CASE__								=   torch.randint(0				,					self.re_embed				,					size=new[unknown].shape		).to(device=new.device		)
									else:
											SCREAMING_SNAKE_CASE__								=   self.unknown_index
									return new.reshape(__A		)


							def   _snake_case	(				self     :Dict				,					__A     :Optional[Any]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   inds.shape
									assert len(__A		) > 1
									SCREAMING_SNAKE_CASE__								=   inds.reshape(ishape[0]				,					-1		)
									SCREAMING_SNAKE_CASE__								=   self.used.to(__A		)
									if self.re_embed > self.used.shape[0]:  # extra token
											SCREAMING_SNAKE_CASE__								=   0  # simply set to zero
									SCREAMING_SNAKE_CASE__								=   torch.gather(used[None, :][inds.shape[0] * [0], :]				,					1				,					__A		)
									return back.reshape(__A		)


							def   _snake_case	(				self     :Optional[int]				,					__A     :List[str]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   z.permute(0				,					2				,					3				,					1		).contiguous()
									SCREAMING_SNAKE_CASE__								=   z.view(-1				,					self.vq_embed_dim		)

									# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
									SCREAMING_SNAKE_CASE__								=   torch.argmin(torch.cdist(__A				,					self.embedding.weight		)				,					dim=1		)

									SCREAMING_SNAKE_CASE__								=   self.embedding(__A		).view(z.shape		)
									SCREAMING_SNAKE_CASE__								=   None
									SCREAMING_SNAKE_CASE__								=   None

									# compute loss for embedding
									if not self.legacy:
											SCREAMING_SNAKE_CASE__								=   self.beta * torch.mean((z_q.detach() - z) ** 2		) + torch.mean((z_q - z.detach()) ** 2		)
									else:
											SCREAMING_SNAKE_CASE__								=   torch.mean((z_q.detach() - z) ** 2		) + self.beta * torch.mean((z_q - z.detach()) ** 2		)

									# preserve gradients
									SCREAMING_SNAKE_CASE__								=   z + (z_q - z).detach()

									# reshape back to match original input shape
									SCREAMING_SNAKE_CASE__								=   z_q.permute(0				,					3				,					1				,					2		).contiguous()

									if self.remap is not None:
											SCREAMING_SNAKE_CASE__								=   min_encoding_indices.reshape(z.shape[0]				,					-1		)  # add batch axis
											SCREAMING_SNAKE_CASE__								=   self.remap_to_used(__A		)
											SCREAMING_SNAKE_CASE__								=   min_encoding_indices.reshape(-1				,					1		)  # flatten

									if self.sane_index_shape:
											SCREAMING_SNAKE_CASE__								=   min_encoding_indices.reshape(z_q.shape[0]				,					z_q.shape[2]				,					z_q.shape[3]		)

									return z_q, loss, (perplexity, min_encodings, min_encoding_indices)


							def   _snake_case	(				self     :Union[str, Any]				,					__A     :Dict				,					__A     :Optional[Any]		) -> Dict:



									"""simple docstring"""



									if self.remap is not None:
											SCREAMING_SNAKE_CASE__								=   indices.reshape(shape[0]				,					-1		)  # add batch axis
											SCREAMING_SNAKE_CASE__								=   self.unmap_to_all(__A		)
											SCREAMING_SNAKE_CASE__								=   indices.reshape(-1		)  # flatten again

									# get quantized latent vectors
									SCREAMING_SNAKE_CASE__								=   self.embedding(__A		)

									if shape is not None:
											SCREAMING_SNAKE_CASE__								=   z_q.view(__A		)
											# reshape back to match original input shape
											SCREAMING_SNAKE_CASE__								=   z_q.permute(0				,					3				,					1				,					2		).contiguous()

									return z_q




class 	UpperCamelCase_							(  UpperCamelCase__     ):


							def __init__(				self     :List[Any]				,					__A     :List[str]				,					__A     :Dict=False		) -> List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   parameters
									SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   torch.chunk(__A				,					2				,					dim=1		)
									SCREAMING_SNAKE_CASE__								=   torch.clamp(self.logvar				,					-3_0.0				,					2_0.0		)
									SCREAMING_SNAKE_CASE__								=   deterministic
									SCREAMING_SNAKE_CASE__								=   torch.exp(0.5 * self.logvar		)
									SCREAMING_SNAKE_CASE__								=   torch.exp(self.logvar		)
									if self.deterministic:
											SCREAMING_SNAKE_CASE__								=   SCREAMING_SNAKE_CASE__								=   torch.zeros_like(
											    self.mean				,					device=self.parameters.device				,					dtype=self.parameters.dtype		)


							def   _snake_case	(				self     :Union[str, Any]				,					__A     :Optional[torch.Generator] = None		) -> torch.FloatTensor:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   randn_tensor(
									    self.mean.shape				,					generator=__A				,					device=self.parameters.device				,					dtype=self.parameters.dtype		)
									SCREAMING_SNAKE_CASE__								=   self.mean + self.std * sample
									return x


							def   _snake_case	(				self     :str				,					__A     :Any=None		) -> Union[str, Any]:



									"""simple docstring"""



									if self.deterministic:
											return torch.Tensor([0.0]		)
									else:
											if other is None:
													return 0.5 * torch.sum(torch.pow(self.mean				,					2		) + self.var - 1.0 - self.logvar				,					dim=[1, 2, 3]		)
											else:
													return 0.5 * torch.sum(
													    torch.pow(self.mean - other.mean				,					2		) / other.var
													    + self.var / other.var
													    - 1.0
													    - self.logvar
													    + other.logvar				,					dim=[1, 2, 3]				,					)


							def   _snake_case	(				self     :str				,					__A     :Optional[int]				,					__A     :List[Any]=[1, 2, 3]		) -> Any:



									"""simple docstring"""



									if self.deterministic:
											return torch.Tensor([0.0]		)
									SCREAMING_SNAKE_CASE__								=   np.log(2.0 * np.pi		)
									return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean				,					2		) / self.var				,					dim=__A		)


							def   _snake_case	(				self     :Any		) -> str:



									"""simple docstring"""



									return self.mean







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[Any]								=pytest.mark.integration






@require_faiss
class 						UpperCAmelCase__ (							__snake_case							):







		def  A__      (					self		):
						_A	:				Any   =      Dataset.from_dict({'''filename''': ['''my_name-train''' + '''_''' + str(A__		) for x in np.arange(30		).tolist()]}		)
						return dset







		def  A__      (					self		):
						import faiss

						_A	:				Dataset   =      self._create_dummy_dataset()
						_A	:				List[str]   =      dset.map(
						    lambda A__       ,A__		: {"vecs": i * np.ones(5       ,dtype=np.floataa		)}       ,with_indices=A__       ,keep_in_memory=A__		)
						_A	:				Tuple   =      dset.add_faiss_index('''vecs'''       ,batch_size=100       ,metric_type=faiss.METRIC_INNER_PRODUCT		)
						_A  ,				_A	:				List[Any]   =      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  A__      (					self		):
						import faiss

						_A	:				Dataset   =      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=100       ,metric_type=faiss.METRIC_INNER_PRODUCT       ,)
						_A  ,				_A	:				Tuple   =      dset.get_nearest_examples('''vecs'''       ,np.ones(5       ,dtype=np.floataa		)		)
						self.assertEqual(examples['''filename'''][0]       ,'''my_name-train_29'''		)







		def  A__      (					self		):
						import faiss

						_A	:				Dataset   =      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=A__		) as tmp_file:
										dset.save_faiss_index('''vecs'''       ,tmp_file.name		)
										dset.load_faiss_index('''vecs2'''       ,tmp_file.name		)
						os.unlink(tmp_file.name		)

						_A  ,				_A	:				Optional[int]   =      dset.get_nearest_examples('''vecs2'''       ,np.ones(5       ,dtype=np.floataa		)		)
						self.assertEqual(examples['''filename'''][0]       ,'''my_name-train_29'''		)







		def  A__      (					self		):
						_A	:				Dataset   =      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(A__       ,partial(dset.get_nearest_examples       ,'''vecs2'''       ,np.ones(5       ,dtype=np.floataa		)		)		)







		def  A__      (					self		):
						from elasticsearch import Elasticsearch

						_A	:				Dataset   =      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:
										_A	:				str   =      {'''acknowledged''': True}
										mocked_bulk.return_value([(True, None)] * 30		)
										_A	:				Optional[Any]   =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 29}]}}
										_A	:				Optional[Any]   =      Elasticsearch()

										dset.add_elasticsearch_index('''filename'''       ,es_client=A__		)
										_A  ,				_A	:				Tuple   =      dset.get_nearest_examples('''filename'''       ,'''my_name-train_29'''		)
										self.assertEqual(examples['''filename'''][0]       ,'''my_name-train_29'''		)






@require_faiss
class 						UpperCAmelCase__ (							__snake_case							):







		def  A__      (					self		):
						import faiss

						_A	:				Union[str, Any]   =      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
						_A	:				str   =      np.zeros(5       ,dtype=np.floataa		)
						_A	:				Dict   =      1
						_A  ,				_A	:				Optional[int]   =      index.search(A__		)
						self.assertRaises(A__       ,index.search       ,query.reshape(-1       ,1		)		)
						self.assertGreater(scores[0]       ,0		)
						self.assertEqual(indices[0]       ,1		)

						# batched queries
						_A	:				Optional[Any]   =      np.eye(5       ,dtype=np.floataa		)[::-1]
						_A  ,				_A	:				Dict   =      index.search_batch(A__		)
						self.assertRaises(A__       ,index.search_batch       ,queries[0]		)
						_A	:				Tuple   =      [scores[0] for scores in total_scores]
						_A	:				List[str]   =      [indices[0] for indices in total_indices]
						self.assertGreater(np.min(A__		)       ,0		)
						self.assertListEqual([4, 3, 2, 1, 0]       ,A__		)







		def  A__      (					self		):
						import faiss

						_A	:				List[Any]   =      FaissIndex(string_factory='''Flat'''		)
						index.add_vectors(np.eye(5       ,dtype=np.floataa		)		)
						self.assertIsInstance(index.faiss_index       ,faiss.IndexFlat		)
						_A	:				str   =      FaissIndex(string_factory='''LSH'''		)
						index.add_vectors(np.eye(5       ,dtype=np.floataa		)		)
						self.assertIsInstance(index.faiss_index       ,faiss.IndexLSH		)
						with self.assertRaises(A__		):
										_A	:				List[str]   =      FaissIndex(string_factory='''Flat'''       ,custom_index=faiss.IndexFlat(5		)		)







		def  A__      (					self		):
						import faiss

						_A	:				Any   =      faiss.IndexFlat(5		)
						_A	:				Any   =      FaissIndex(custom_index=A__		)
						index.add_vectors(np.eye(5       ,dtype=np.floataa		)		)
						self.assertIsInstance(index.faiss_index       ,faiss.IndexFlat		)







		def  A__      (					self		):
						import faiss

						_A	:				Optional[int]   =      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=A__		) as tmp_file:
										index.save(tmp_file.name		)
										_A	:				str   =      FaissIndex.load(tmp_file.name		)
						os.unlink(tmp_file.name		)

						_A	:				Tuple   =      np.zeros(5       ,dtype=np.floataa		)
						_A	:				Tuple   =      1
						_A  ,				_A	:				Union[str, Any]   =      index.search(A__		)
						self.assertGreater(scores[0]       ,0		)
						self.assertEqual(indices[0]       ,1		)



@require_faiss
def  a__ (__lowercase							:List[Any]   )		-> int:
				import faiss

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

				_A	:				List[Any]   =      '''index.faiss'''
				_A	:				Any   =      f"""mock://{index_name}"""
				index.save(__lowercase					,  storage_options=mockfs.storage_options   )
				_A	:				Dict   =      FaissIndex.load(__lowercase					,  storage_options=mockfs.storage_options   )

				_A	:				Tuple   =      np.zeros(5					,  dtype=np.floataa   )
				_A	:				Union[str, Any]   =      1
				_A  ,				_A	:				Optional[int]   =      index.search(__lowercase   )
				assert scores[0] > 0
				assert indices[0] == 1






@require_elasticsearch
class 						UpperCAmelCase__ (							__snake_case							):







		def  A__      (					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:
										_A	:				Tuple   =      Elasticsearch()
										_A	:				Tuple   =      {'''acknowledged''': True}
										_A	:				int   =      ElasticSearchIndex(es_client=A__		)
										mocked_bulk.return_value([(True, None)] * 3		)
										index.add_documents(['''foo''', '''bar''', '''foobar''']		)

										# single query
										_A	:				Dict   =      '''foo'''
										_A	:				Optional[Any]   =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 0}]}}
										_A  ,				_A	:				Optional[Any]   =      index.search(A__		)
										self.assertEqual(scores[0]       ,1		)
										self.assertEqual(indices[0]       ,0		)

										# single query with timeout
										_A	:				Union[str, Any]   =      '''foo'''
										_A	:				Optional[int]   =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 0}]}}
										_A  ,				_A	:				Optional[int]   =      index.search(A__       ,request_timeout=30		)
										self.assertEqual(scores[0]       ,1		)
										self.assertEqual(indices[0]       ,0		)

										# batched queries
										_A	:				Dict   =      ['''foo''', '''bar''', '''foobar''']
										_A	:				Dict   =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
										_A  ,				_A	:				Dict   =      index.search_batch(A__		)
										_A	:				Union[str, Any]   =      [scores[0] for scores in total_scores]
										_A	:				Tuple   =      [indices[0] for indices in total_indices]
										self.assertGreater(np.min(A__		)       ,0		)
										self.assertListEqual([1, 1, 1]       ,A__		)

										# batched queries with timeout
										_A	:				Union[str, Any]   =      ['''foo''', '''bar''', '''foobar''']
										_A	:				int   =      {'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
										_A  ,				_A	:				str   =      index.search_batch(A__       ,request_timeout=30		)
										_A	:				Dict   =      [scores[0] for scores in total_scores]
										_A	:				List[str]   =      [indices[0] for indices in total_indices]
										self.assertGreater(np.min(A__		)       ,0		)
										self.assertListEqual([1, 1, 1]       ,A__		)



