





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







@require_faiss
class   _UpperCamelCase  (       __lowercase  ):







     '''simple docstring'''

     def        SCREAMING_SNAKE_CASE							(						self    :  Dict		):






      '''simple docstring'''







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





     def        SCREAMING_SNAKE_CASE							(						self    :  Tuple		):






      '''simple docstring'''







      import faiss

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





     def        SCREAMING_SNAKE_CASE							(						self    :  int		):






      '''simple docstring'''







      import faiss

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





     def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






      '''simple docstring'''







      import faiss

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

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





     def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any]		):






      '''simple docstring'''







      UpperCAmelCase_								=     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(__lowercase ,					partial(dset.get_nearest_examples ,					"""vecs2""" ,					np.ones(5 ,					dtype=np.floataa		)		)		)





     def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any]		):






      '''simple docstring'''







      from elasticsearch import Elasticsearch

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

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







@require_faiss
class   _UpperCamelCase  (       __lowercase  ):







     '''simple docstring'''

     def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int]		):






      '''simple docstring'''







      import faiss

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

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





     def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






      '''simple docstring'''







      import faiss

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





     def        SCREAMING_SNAKE_CASE							(						self    :  List[str]		):






      '''simple docstring'''







      import faiss

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





     def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






      '''simple docstring'''







      import faiss

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

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




@require_faiss
def    A_(			A   ):
 import faiss

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

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

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







@require_elasticsearch
class   _UpperCamelCase  (       __lowercase  ):







     '''simple docstring'''

     def        SCREAMING_SNAKE_CASE							(						self    :  List[str]		):






      '''simple docstring'''







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

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

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

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

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









from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import (
    BaseOutput,
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_onnx_available,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)







@dataclass
class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					lowerCamelCase     :     Union[List[PIL.Image.Image], np.ndarray]
					lowerCamelCase     :     Optional[List[bool]]


try:
				if not (is_transformers_available() and is_torch_available()):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
				from .pipeline_cycle_diffusion import CycleDiffusionPipeline
				from .pipeline_stable_diffusion import StableDiffusionPipeline
				from .pipeline_stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline
				from .pipeline_stable_diffusion_imgaimg import StableDiffusionImgaImgPipeline
				from .pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline
				from .pipeline_stable_diffusion_inpaint_legacy import StableDiffusionInpaintPipelineLegacy
				from .pipeline_stable_diffusion_instruct_pixapix import StableDiffusionInstructPixaPixPipeline
				from .pipeline_stable_diffusion_latent_upscale import StableDiffusionLatentUpscalePipeline
				from .pipeline_stable_diffusion_ldmad import StableDiffusionLDMaDPipeline
				from .pipeline_stable_diffusion_model_editing import StableDiffusionModelEditingPipeline
				from .pipeline_stable_diffusion_panorama import StableDiffusionPanoramaPipeline
				from .pipeline_stable_diffusion_paradigms import StableDiffusionParadigmsPipeline
				from .pipeline_stable_diffusion_sag import StableDiffusionSAGPipeline
				from .pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
				from .pipeline_stable_unclip import StableUnCLIPPipeline
				from .pipeline_stable_unclip_imgaimg import StableUnCLIPImgaImgPipeline
				from .safety_checker import StableDiffusionSafetyChecker
				from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer

try:
				if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import StableDiffusionImageVariationPipeline
else:
				from .pipeline_stable_diffusion_image_variation import StableDiffusionImageVariationPipeline


try:
				if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.26.0""")):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import (
				    StableDiffusionDepthaImgPipeline,
				    StableDiffusionDiffEditPipeline,
				    StableDiffusionPixaPixZeroPipeline,
				)
else:
				from .pipeline_stable_diffusion_depthaimg import StableDiffusionDepthaImgPipeline
				from .pipeline_stable_diffusion_diffedit import StableDiffusionDiffEditPipeline
				from .pipeline_stable_diffusion_pixapix_zero import StableDiffusionPixaPixZeroPipeline


try:
				if not (
				    is_torch_available()
				    and is_transformers_available()
				    and is_k_diffusion_available()
				    and is_k_diffusion_version(""">=""", """0.0.12""")
				):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
				from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline

try:
				if not (is_transformers_available() and is_onnx_available()):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_onnx_objects import *  # noqa F403
else:
				from .pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline, StableDiffusionOnnxPipeline
				from .pipeline_onnx_stable_diffusion_imgaimg import OnnxStableDiffusionImgaImgPipeline
				from .pipeline_onnx_stable_diffusion_inpaint import OnnxStableDiffusionInpaintPipeline
				from .pipeline_onnx_stable_diffusion_inpaint_legacy import OnnxStableDiffusionInpaintPipelineLegacy
				from .pipeline_onnx_stable_diffusion_upscale import OnnxStableDiffusionUpscalePipeline

if is_transformers_available() and is_flax_available():
				import flax







				@flax.struct.dataclass
				class   _UpperCamelCase  (       A_  ):







									'''simple docstring'''

									lowerCamelCase     :     np.ndarray
									lowerCamelCase     :     List[bool]


				from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState
				from .pipeline_flax_stable_diffusion import FlaxStableDiffusionPipeline
				from .pipeline_flax_stable_diffusion_imgaimg import FlaxStableDiffusionImgaImgPipeline
				from .pipeline_flax_stable_diffusion_inpaint import FlaxStableDiffusionInpaintPipeline
				from .safety_checker_flax import FlaxStableDiffusionSafetyChecker


