




'''simple docstring'''


import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class 						a_				(			UpperCAmelCase__     ,						unittest.TestCase						):
	lowercase_       :     Optional[Any]			  =      AudioLDMPipeline
	lowercase_       :     List[str]			  =      TEXT_TO_AUDIO_PARAMS
	lowercase_       :     List[Any]			  =      TEXT_TO_AUDIO_BATCH_PARAMS
	lowercase_       :     Dict			  =      frozenset(
	    [
	        '''num_inference_steps''',
	        '''num_waveforms_per_prompt''',
	        '''generator''',
	        '''latents''',
	        '''output_type''',
	        '''return_dict''',
	        '''callback''',
	        '''callback_steps''',
	    ]						)




	def   lowercase__   (       self     :  int      ):
							torch.manual_seed(0      )
							__snake_case						 =   UNetaDConditionModel(
							    block_out_channels=(3_2, 6_4)       ,		layers_per_block=2       ,		sample_size=3_2       ,		in_channels=4       ,		out_channels=4       ,		down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')       ,		up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')       ,		cross_attention_dim=(3_2, 6_4)       ,		class_embed_type='simple_projection'       ,		projection_class_embeddings_input_dim=3_2       ,		class_embeddings_concat=__lowerCAmelCase       ,		)
							__snake_case						 =   DDIMScheduler(
							    beta_start=0.00085       ,		beta_end=0.012       ,		beta_schedule='scaled_linear'       ,		clip_sample=__lowerCAmelCase       ,		set_alpha_to_one=__lowerCAmelCase       ,		)
							torch.manual_seed(0      )
							__snake_case						 =   AutoencoderKL(
							    block_out_channels=[3_2, 6_4]       ,		in_channels=1       ,		out_channels=1       ,		down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']       ,		up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']       ,		latent_channels=4       ,		)
							torch.manual_seed(0      )
							__snake_case						 =   ClapTextConfig(
							    bos_token_id=0       ,		eos_token_id=2       ,		hidden_size=3_2       ,		intermediate_size=3_7       ,		layer_norm_eps=1E-05       ,		num_attention_heads=4       ,		num_hidden_layers=5       ,		pad_token_id=1       ,		vocab_size=1_0_0_0       ,		projection_dim=3_2       ,		)
							__snake_case						 =   ClapTextModelWithProjection(__lowerCAmelCase      )
							__snake_case						 =   RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'       ,		model_max_length=7_7      )

							__snake_case						 =   SpeechTaHifiGanConfig(
							    model_in_dim=8       ,		sampling_rate=1_6_0_0_0       ,		upsample_initial_channel=1_6       ,		upsample_rates=[2, 2]       ,		upsample_kernel_sizes=[4, 4]       ,		resblock_kernel_sizes=[3, 7]       ,		resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]       ,		normalize_before=__lowerCAmelCase       ,		)

							__snake_case						 =   SpeechTaHifiGan(__lowerCAmelCase      )

							__snake_case						 =   {
							    'unet': unet,
							    'scheduler': scheduler,
							    'vae': vae,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'vocoder': vocoder,
							}
							return components




	def   lowercase__   (       self     :  Any       ,		__lowerCAmelCase     :  List[str]       ,		__lowerCAmelCase     :  List[Any]=0      ):
							if str(__lowerCAmelCase      ).startswith('mps'      ):
													__snake_case						 =   torch.manual_seed(__lowerCAmelCase      )
							else:
													__snake_case						 =   torch.Generator(device=__lowerCAmelCase      ).manual_seed(__lowerCAmelCase      )
							__snake_case						 =   {
							    'prompt': 'A hammer hitting a wooden surface',
							    'generator': generator,
							    'num_inference_steps': 2,
							    'guidance_scale': 6.0,
							}
							return inputs




	def   lowercase__   (       self     :  str      ):
							__snake_case						 =   'cpu'  # ensure determinism for the device-dependent torch.Generator

							__snake_case						 =   self.get_dummy_components()
							__snake_case						 =   AudioLDMPipeline(**__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

							__snake_case						 =   self.get_dummy_inputs(__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe(**__lowerCAmelCase      )
							__snake_case						 =   output.audios[0]

							assert audio.ndim == 1
							assert len(__lowerCAmelCase      ) == 2_5_6

							__snake_case						 =   audio[:1_0]
							__snake_case						 =   np.array(
							    [-0.0050, 0.0050, -0.0060, 0.0033, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0033]      )

							assert np.abs(audio_slice - expected_slice      ).max() < 1E-2




	def   lowercase__   (       self     :  List[Any]      ):
							__snake_case						 =   self.get_dummy_components()
							__snake_case						 =   AudioLDMPipeline(**__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

							__snake_case						 =   self.get_dummy_inputs(__lowerCAmelCase      )
							__snake_case						 =   3 * [inputs['prompt']]

							# forward
							__snake_case						 =   audioldm_pipe(**__lowerCAmelCase      )
							__snake_case						 =   output.audios[0]

							__snake_case						 =   self.get_dummy_inputs(__lowerCAmelCase      )
							__snake_case						 =   3 * [inputs.pop('prompt'      )]

							__snake_case						 =   audioldm_pipe.tokenizer(
							    __lowerCAmelCase       ,		padding='max_length'       ,		max_length=audioldm_pipe.tokenizer.model_max_length       ,		truncation=__lowerCAmelCase       ,		return_tensors='pt'       ,		)
							__snake_case						 =   text_inputs['input_ids'].to(__lowerCAmelCase      )

							__snake_case						 =   audioldm_pipe.text_encoder(
							    __lowerCAmelCase       ,		)
							__snake_case						 =   prompt_embeds.text_embeds
							# additional L_2 normalization over each hidden-state
							__snake_case						 =   F.normalize(__lowerCAmelCase       ,		dim=-1      )

							__snake_case						 =   prompt_embeds

							# forward
							__snake_case						 =   audioldm_pipe(**__lowerCAmelCase      )
							__snake_case						 =   output.audios[0]

							assert np.abs(audio_a - audio_a      ).max() < 1E-2




	def   lowercase__   (       self     :  Optional[int]      ):
							__snake_case						 =   self.get_dummy_components()
							__snake_case						 =   AudioLDMPipeline(**__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

							__snake_case						 =   self.get_dummy_inputs(__lowerCAmelCase      )
							__snake_case						 =   3 * ['this is a negative prompt']
							__snake_case						 =   negative_prompt
							__snake_case						 =   3 * [inputs['prompt']]

							# forward
							__snake_case						 =   audioldm_pipe(**__lowerCAmelCase      )
							__snake_case						 =   output.audios[0]

							__snake_case						 =   self.get_dummy_inputs(__lowerCAmelCase      )
							__snake_case						 =   3 * [inputs.pop('prompt'      )]

							__snake_case						 =   []
							for p in [prompt, negative_prompt]:
													__snake_case						 =   audioldm_pipe.tokenizer(
													    __lowerCAmelCase       ,		padding='max_length'       ,		max_length=audioldm_pipe.tokenizer.model_max_length       ,		truncation=__lowerCAmelCase       ,		return_tensors='pt'       ,		)
													__snake_case						 =   text_inputs['input_ids'].to(__lowerCAmelCase      )

													__snake_case						 =   audioldm_pipe.text_encoder(
													    __lowerCAmelCase       ,		)
													__snake_case						 =   text_embeds.text_embeds
													# additional L_2 normalization over each hidden-state
													__snake_case						 =   F.normalize(__lowerCAmelCase       ,		dim=-1      )

													embeds.append(__lowerCAmelCase      )

							__snake_case ,       __snake_case						 =   embeds

							# forward
							__snake_case						 =   audioldm_pipe(**__lowerCAmelCase      )
							__snake_case						 =   output.audios[0]

							assert np.abs(audio_a - audio_a      ).max() < 1E-2




	def   lowercase__   (       self     :  List[str]      ):
							__snake_case						 =   'cpu'  # ensure determinism for the device-dependent torch.Generator
							__snake_case						 =   self.get_dummy_components()
							__snake_case						 =   PNDMScheduler(skip_prk_steps=__lowerCAmelCase      )
							__snake_case						 =   AudioLDMPipeline(**__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

							__snake_case						 =   self.get_dummy_inputs(__lowerCAmelCase      )
							__snake_case						 =   'egg cracking'
							__snake_case						 =   audioldm_pipe(**__lowerCAmelCase       ,		negative_prompt=__lowerCAmelCase      )
							__snake_case						 =   output.audios[0]

							assert audio.ndim == 1
							assert len(__lowerCAmelCase      ) == 2_5_6

							__snake_case						 =   audio[:1_0]
							__snake_case						 =   np.array(
							    [-0.0051, 0.0050, -0.0060, 0.0034, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0032]      )

							assert np.abs(audio_slice - expected_slice      ).max() < 1E-2




	def   lowercase__   (       self     :  Optional[int]      ):
							__snake_case						 =   'cpu'  # ensure determinism for the device-dependent torch.Generator
							__snake_case						 =   self.get_dummy_components()
							__snake_case						 =   PNDMScheduler(skip_prk_steps=__lowerCAmelCase      )
							__snake_case						 =   AudioLDMPipeline(**__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

							__snake_case						 =   'A hammer hitting a wooden surface'

							# test num_waveforms_per_prompt=1 (default)
							__snake_case						 =   audioldm_pipe(__lowerCAmelCase       ,		num_inference_steps=2      ).audios

							assert audios.shape == (1, 2_5_6)

							# test num_waveforms_per_prompt=1 (default) for batch of prompts
							__snake_case						 =   2
							__snake_case						 =   audioldm_pipe([prompt] * batch_size       ,		num_inference_steps=2      ).audios

							assert audios.shape == (batch_size, 2_5_6)

							# test num_waveforms_per_prompt for single prompt
							__snake_case						 =   2
							__snake_case						 =   audioldm_pipe(__lowerCAmelCase       ,		num_inference_steps=2       ,		num_waveforms_per_prompt=__lowerCAmelCase      ).audios

							assert audios.shape == (num_waveforms_per_prompt, 2_5_6)

							# test num_waveforms_per_prompt for batch of prompts
							__snake_case						 =   2
							__snake_case						 =   audioldm_pipe(
							    [prompt] * batch_size       ,		num_inference_steps=2       ,		num_waveforms_per_prompt=__lowerCAmelCase      ).audios

							assert audios.shape == (batch_size * num_waveforms_per_prompt, 2_5_6)




	def   lowercase__   (       self     :  str      ):
							__snake_case						 =   'cpu'  # ensure determinism for the device-dependent torch.Generator
							__snake_case						 =   self.get_dummy_components()
							__snake_case						 =   AudioLDMPipeline(**__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.vocoder.config.sampling_rate

							__snake_case						 =   self.get_dummy_inputs(__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe(audio_length_in_s=0.016       ,		**__lowerCAmelCase      )
							__snake_case						 =   output.audios[0]

							assert audio.ndim == 1
							assert len(__lowerCAmelCase      ) / vocoder_sampling_rate == 0.016

							__snake_case						 =   audioldm_pipe(audio_length_in_s=0.032       ,		**__lowerCAmelCase      )
							__snake_case						 =   output.audios[0]

							assert audio.ndim == 1
							assert len(__lowerCAmelCase      ) / vocoder_sampling_rate == 0.032




	def   lowercase__   (       self     :  List[str]      ):
							__snake_case						 =   self.get_dummy_components()
							__snake_case						 =   AudioLDMPipeline(**__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

							__snake_case						 =   ['hey']

							__snake_case						 =   audioldm_pipe(__lowerCAmelCase       ,		num_inference_steps=1      )
							__snake_case						 =   output.audios.shape
							assert audio_shape == (1, 2_5_6)

							__snake_case						 =   audioldm_pipe.vocoder.config
							config.model_in_dim *= 2
							__snake_case						 =   SpeechTaHifiGan(__lowerCAmelCase      ).to(__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe(__lowerCAmelCase       ,		num_inference_steps=1      )
							__snake_case						 =   output.audios.shape
							# waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
							assert audio_shape == (1, 2_5_6)




	def   lowercase__   (       self     :  Union[str, Any]      ):
							self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__lowerCAmelCase      )




	def   lowercase__   (       self     :  Any      ):
							self._test_inference_batch_single_identical(test_mean_pixel_difference=__lowerCAmelCase      )




	@unittest.skipIf(
	    torch_device != 'cuda' or not is_xformers_available()       ,		reason='XFormers attention is only available with CUDA and `xformers` installed'       ,		)
	def   lowercase__   (       self     :  List[str]      ):
							self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__lowerCAmelCase      )





@slow
class 						a_				(			unittest.TestCase						):




	def   lowercase__   (       self     :  List[Any]      ):
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




	def   lowercase__   (       self     :  Dict       ,		__lowerCAmelCase     :  int       ,		__lowerCAmelCase     :  Dict="cpu"       ,		__lowerCAmelCase     :  List[Any]=torch.floataa       ,		__lowerCAmelCase     :  int=0      ):
							__snake_case						 =   torch.Generator(device=__lowerCAmelCase      ).manual_seed(__lowerCAmelCase      )
							__snake_case						 =   np.random.RandomState(__lowerCAmelCase      ).standard_normal((1, 8, 1_2_8, 1_6)      )
							__snake_case						 =   torch.from_numpy(__lowerCAmelCase      ).to(device=__lowerCAmelCase       ,		dtype=__lowerCAmelCase      )
							__snake_case						 =   {
							    'prompt': 'A hammer hitting a wooden surface',
							    'latents': latents,
							    'generator': generator,
							    'num_inference_steps': 3,
							    'guidance_scale': 2.5,
							}
							return inputs




	def   lowercase__   (       self     :  Tuple      ):
							__snake_case						 =   AudioLDMPipeline.from_pretrained('cvssp/audioldm'      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

							__snake_case						 =   self.get_inputs(__lowerCAmelCase      )
							__snake_case						 =   2_5
							__snake_case						 =   audioldm_pipe(**__lowerCAmelCase      ).audios[0]

							assert audio.ndim == 1
							assert len(__lowerCAmelCase      ) == 8_1_9_2_0

							__snake_case						 =   audio[7_7_2_3_0:7_7_2_4_0]
							__snake_case						 =   np.array(
							    [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315]      )
							__snake_case						 =   np.abs(expected_slice - audio_slice      ).max()
							assert max_diff < 1E-2




	def   lowercase__   (       self     :  Union[str, Any]      ):
							__snake_case						 =   AudioLDMPipeline.from_pretrained('cvssp/audioldm'      )
							__snake_case						 =   LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config      )
							__snake_case						 =   audioldm_pipe.to(__lowerCAmelCase      )
							audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase      )

							__snake_case						 =   self.get_inputs(__lowerCAmelCase      )
							__snake_case						 =   audioldm_pipe(**__lowerCAmelCase      ).audios[0]

							assert audio.ndim == 1
							assert len(__lowerCAmelCase      ) == 8_1_9_2_0

							__snake_case						 =   audio[2_7_7_8_0:2_7_7_9_0]
							__snake_case						 =   np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212]      )
							__snake_case						 =   np.abs(expected_slice - audio_slice      ).max()
							assert max_diff < 3E-2







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


_lowercase  =			pytest.mark.integration
@require_faiss
class 						a_				(			UpperCAmelCase__						):




	def   lowercase__   (       self     :  List[Any]      ):
							__snake_case						 =   Dataset.from_dict({'filename': ['my_name-train' + '_' + str(__lowerCAmelCase      ) for x in np.arange(3_0      ).tolist()]}      )
							return dset




	def   lowercase__   (       self     :  List[str]      ):
							import faiss

							__snake_case						 =   self._create_dummy_dataset()
							__snake_case						 =   dset.map(
							    lambda __lowerCAmelCase       ,		__lowerCAmelCase      : {"vecs": i * np.ones(5       ,		dtype=np.floataa      )}       ,		with_indices=__lowerCAmelCase       ,		keep_in_memory=__lowerCAmelCase      )
							__snake_case						 =   dset.add_faiss_index('vecs'       ,		batch_size=1_0_0       ,		metric_type=faiss.METRIC_INNER_PRODUCT      )
							__snake_case ,       __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   lowercase__   (       self     :  Optional[int]      ):
							import faiss

							__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       ,		)
							__snake_case ,       __snake_case						 =   dset.get_nearest_examples('vecs'       ,		np.ones(5       ,		dtype=np.floataa      )      )
							self.assertEqual(examples['filename'][0]       ,		'my_name-train_29'      )




	def   lowercase__   (       self     :  Dict      ):
							import faiss

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

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




	def   lowercase__   (       self     :  Union[str, Any]      ):
							__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(__lowerCAmelCase       ,		partial(dset.get_nearest_examples       ,		'vecs2'       ,		np.ones(5       ,		dtype=np.floataa      )      )      )




	def   lowercase__   (       self     :  List[str]      ):
							from elasticsearch import Elasticsearch

							__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:
													__snake_case						 =   {'acknowledged': True}
													mocked_bulk.return_value([(True, None)] * 3_0      )
													__snake_case						 =   {'hits': {'hits': [{'_score': 1, '_id': 2_9}]}}
													__snake_case						 =   Elasticsearch()

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





@require_faiss
class 						a_				(			UpperCAmelCase__						):




	def   lowercase__   (       self     :  Dict      ):
							import faiss

							__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
							__snake_case						 =   np.zeros(5       ,		dtype=np.floataa      )
							__snake_case						 =   1
							__snake_case ,       __snake_case						 =   index.search(__lowerCAmelCase      )
							self.assertRaises(__lowerCAmelCase       ,		index.search       ,		query.reshape(-1       ,		1      )      )
							self.assertGreater(scores[0]       ,		0      )
							self.assertEqual(indices[0]       ,		1      )

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




	def   lowercase__   (       self     :  Optional[int]      ):
							import faiss

							__snake_case						 =   FaissIndex(string_factory='Flat'      )
							index.add_vectors(np.eye(5       ,		dtype=np.floataa      )      )
							self.assertIsInstance(index.faiss_index       ,		faiss.IndexFlat      )
							__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(__lowerCAmelCase      ):
													__snake_case						 =   FaissIndex(string_factory='Flat'       ,		custom_index=faiss.IndexFlat(5      )      )




	def   lowercase__   (       self     :  int      ):
							import faiss

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




	def   lowercase__   (       self     :  Tuple      ):
							import faiss

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

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





@require_faiss
def 						lowerCamelCase__   (				a							):
						import faiss

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

						__snake_case						 =   'index.faiss'
						__snake_case						 =   f'mock://{index_name}'
						index.save(a    ,	storage_options=mockfs.storage_options							)
						__snake_case						 =   FaissIndex.load(a    ,	storage_options=mockfs.storage_options							)

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






@require_elasticsearch
class 						a_				(			UpperCAmelCase__						):




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

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

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

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

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

