


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







@require_faiss
class lowerCAmelCase_						(			lowerCamelCase__	):

	'''simple docstring'''

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



	def 						UpperCamelCase__					(     self ):
							import faiss

							snake_case_	=				self._create_dummy_dataset()
							snake_case_	=				dset.map(
							    lambda _UpperCAmelCase		,				_UpperCAmelCase : {"vecs": i * np.ones(5		,				dtype=np.floataa )}		,				with_indices=_UpperCAmelCase		,				keep_in_memory=_UpperCAmelCase )
							snake_case_	=				dset.add_faiss_index('''vecs'''		,				batch_size=1_00		,				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 						UpperCamelCase__					(     self ):
							import faiss

							snake_case_	=				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		,				)
							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 						UpperCamelCase__					(     self ):
							import faiss

							snake_case_	=				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=_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 )

							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 						UpperCamelCase__					(     self ):
							snake_case_	=				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(_UpperCAmelCase		,				partial(dset.get_nearest_examples		,				'''vecs2'''		,				np.ones(5		,				dtype=np.floataa ) ) )





	def 						UpperCamelCase__					(     self ):
							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)] * 30 )
													snake_case_	=				{'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 29}]}}
													snake_case_	=				Elasticsearch()

													dset.add_elasticsearch_index('''filename'''		,				es_client=_UpperCAmelCase )
													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 lowerCAmelCase_						(			lowerCamelCase__	):

	'''simple docstring'''

	def 						UpperCamelCase__					(     self ):
							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		,				10 )

							# single query
							snake_case_	=				np.zeros(5		,				dtype=np.floataa )
							snake_case_	=				1
							snake_case_	,   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
							snake_case_	=				np.eye(5		,				dtype=np.floataa )[::-1]
							snake_case_	,   snake_case_	=				index.search_batch(_UpperCAmelCase )
							self.assertRaises(_UpperCAmelCase		,				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(_UpperCAmelCase )		,				0 )
							self.assertListEqual([4, 3, 2, 1, 0]		,				_UpperCAmelCase )



	def 						UpperCamelCase__					(     self ):
							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(_UpperCAmelCase ):
													snake_case_	=				FaissIndex(string_factory='''Flat'''		,				custom_index=faiss.IndexFlat(5 ) )



	def 						UpperCamelCase__					(     self ):
							import faiss

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





	def 						UpperCamelCase__					(     self ):
							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=_UpperCAmelCase ) 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(_UpperCAmelCase )
							self.assertGreater(scores[0]		,				0 )
							self.assertEqual(indices[0]		,				1 )






@require_faiss
def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  Any:




						"""simple docstring"""

						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(SCREAMING_SNAKE_CASE , storage_options=mockfs.storage_options )
						snake_case_	=				FaissIndex.load(SCREAMING_SNAKE_CASE , storage_options=mockfs.storage_options )

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







@require_elasticsearch
class lowerCAmelCase_						(			lowerCamelCase__	):

	'''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:
													snake_case_	=				Elasticsearch()
													snake_case_	=				{'''acknowledged''': True}
													snake_case_	=				ElasticSearchIndex(es_client=_UpperCAmelCase )
													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(_UpperCAmelCase )
													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(_UpperCAmelCase		,				request_timeout=30 )
													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(_UpperCAmelCase )
													snake_case_	=				[scores[0] for scores in total_scores]
													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
													snake_case_	=				['''foo''', '''bar''', '''foobar''']
													snake_case_	=				{'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
													snake_case_	,   snake_case_	=				index.search_batch(_UpperCAmelCase		,				request_timeout=30 )
													snake_case_	=				[scores[0] for scores in total_scores]
													snake_case_	=				[indices[0] for indices in total_indices]
													self.assertGreater(np.min(_UpperCAmelCase )		,				0 )
													self.assertListEqual([1, 1, 1]		,				_UpperCAmelCase )



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


UpperCAmelCase      =   False







class lowerCAmelCase_						(			unittest.TestCase	):

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 @property
 def 						UpperCamelCase__					(     self ):
       return 12



 @property
 def 						UpperCamelCase__					(     self ):
       return 12



 @property
 def 						UpperCamelCase__					(     self ):
       return 32



 @property
 def 						UpperCamelCase__					(     self ):
       torch.manual_seed(0 )
       snake_case_	=				VQModel(
           block_out_channels=[32, 64]		,				in_channels=3		,				out_channels=3		,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']		,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']		,				latent_channels=3		,				num_vq_embeddings=self.num_embed		,				vq_embed_dim=3		,				)
       return model



 @property
 def 						UpperCamelCase__					(     self ):
       snake_case_	=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )
       return tokenizer



 @property
 def 						UpperCamelCase__					(     self ):
       torch.manual_seed(0 )
       snake_case_	=				CLIPTextConfig(
           bos_token_id=0		,				eos_token_id=2		,				hidden_size=self.text_embedder_hidden_size		,				intermediate_size=37		,				layer_norm_eps=1E-05		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=10_00		,				)
       return CLIPTextModel(_UpperCAmelCase )



 @property
 def 						UpperCamelCase__					(     self ):
       torch.manual_seed(0 )

       snake_case_	=				12
       snake_case_	=				12

       snake_case_	=				{
           '''attention_bias''': True,
           '''cross_attention_dim''': 32,
           '''attention_head_dim''': height * width,
           '''num_attention_heads''': 1,
           '''num_vector_embeds''': self.num_embed,
           '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
           '''norm_num_groups''': 32,
           '''sample_size''': width,
           '''activation_fn''': '''geglu-approximate''',
       }

       snake_case_	=				TransformeraDModel(**_UpperCAmelCase )
       return model



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				'''cpu'''

       snake_case_	=				self.dummy_vqvae
       snake_case_	=				self.dummy_text_encoder
       snake_case_	=				self.dummy_tokenizer
       snake_case_	=				self.dummy_transformer
       snake_case_	=				VQDiffusionScheduler(self.num_embed )
       snake_case_	=				LearnedClassifierFreeSamplingEmbeddings(learnable=_UpperCAmelCase )

       snake_case_	=				VQDiffusionPipeline(
           vqvae=_UpperCAmelCase		,				text_encoder=_UpperCAmelCase		,				tokenizer=_UpperCAmelCase		,				transformer=_UpperCAmelCase		,				scheduler=_UpperCAmelCase		,				learned_classifier_free_sampling_embeddings=_UpperCAmelCase		,				)
       snake_case_	=				pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				'''teddy bear playing in the pool'''

       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipe([prompt]		,				generator=_UpperCAmelCase		,				num_inference_steps=2		,				output_type='''np''' )
       snake_case_	=				output.images

       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipe(
           [prompt]		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				return_dict=_UpperCAmelCase		,				num_inference_steps=2 )[0]

       snake_case_	=				image[0, -3:, -3:, -1]
       snake_case_	=				image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       snake_case_	=				np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				'''cpu'''

       snake_case_	=				self.dummy_vqvae
       snake_case_	=				self.dummy_text_encoder
       snake_case_	=				self.dummy_tokenizer
       snake_case_	=				self.dummy_transformer
       snake_case_	=				VQDiffusionScheduler(self.num_embed )
       snake_case_	=				LearnedClassifierFreeSamplingEmbeddings(
           learnable=_UpperCAmelCase		,				hidden_size=self.text_embedder_hidden_size		,				length=tokenizer.model_max_length )

       snake_case_	=				VQDiffusionPipeline(
           vqvae=_UpperCAmelCase		,				text_encoder=_UpperCAmelCase		,				tokenizer=_UpperCAmelCase		,				transformer=_UpperCAmelCase		,				scheduler=_UpperCAmelCase		,				learned_classifier_free_sampling_embeddings=_UpperCAmelCase		,				)
       snake_case_	=				pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				'''teddy bear playing in the pool'''

       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipe([prompt]		,				generator=_UpperCAmelCase		,				num_inference_steps=2		,				output_type='''np''' )
       snake_case_	=				output.images

       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipe(
           [prompt]		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				return_dict=_UpperCAmelCase		,				num_inference_steps=2 )[0]

       snake_case_	=				image[0, -3:, -3:, -1]
       snake_case_	=				image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       snake_case_	=				np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 2.0
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2







@slow
@require_torch_gpu
class lowerCAmelCase_						(			unittest.TestCase	):

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy''' )

       snake_case_	=				VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq''' )
       snake_case_	=				pipeline.to(_UpperCAmelCase )
       pipeline.set_progress_bar_config(disable=_UpperCAmelCase )

       # requires GPU generator for gumbel softmax
       # don't use GPU generator in tests though
       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipeline(
           '''teddy bear playing in the pool'''		,				num_images_per_prompt=1		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				)

       snake_case_	=				output.images[0]

       assert image.shape == (2_56, 2_56, 3)
       assert np.abs(expected_image - image ).max() < 2.0