


import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

  if is_vision_available():
    from transformers import MaskaFormerImageProcessor

if is_vision_available():
  from PIL import Image


class     __UpperCAmelCase							:


     """simple docstring"""


     def __init__(	self  ,       __A  ,       __A=2  ,       __A=True  ,       __A=False  ,       __A=10  ,       __A=3  ,       __A=32 * 8  ,       __A=32 * 8  ,       __A=4  ,       __A=64  ,       ):
        __a					=  parent
        __a					=  batch_size
        __a					=  is_training
        __a					=  use_auxiliary_loss
        __a					=  num_queries
        __a					=  num_channels
        __a					=  min_size
        __a					=  max_size
        __a					=  num_labels
        __a					=  hidden_dim
        __a					=  hidden_dim


     def 	snake_case_  (	self    ):
        __a					=  floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]    ).to(
            __A    )

        __a					=  torch.ones([self.batch_size, self.min_size, self.max_size]  ,       device=__A    )

        __a					=  (
            torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]  ,       device=__A    ) > 0.5
        ).float()
        __a					=  (torch.rand((self.batch_size, self.num_labels)  ,       device=__A    ) > 0.5).long()

        __a					=  self.get_config()
        return config, pixel_values, pixel_mask, mask_labels, class_labels


     def 	snake_case_  (	self    ):
        __a					=  MaskaFormerConfig(
            hidden_size=self.hidden_dim  ,       )
        __a					=  self.num_queries
        __a					=  self.num_labels

        __a					=  [1, 1, 1, 1]
        __a					=  self.num_channels

        __a					=  64
        __a					=  128
        __a					=  self.hidden_dim
        __a					=  self.hidden_dim
        __a					=  self.hidden_dim
        return config


     def 	snake_case_  (	self    ):
        __a	,    __a	,    __a	,    __a	,    __a					=  self.prepare_config_and_inputs()
        __a					=  {"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
        return config, inputs_dict


     def 	snake_case_  (	self  ,       __A  ,       __A    ):
        __a					=  output.encoder_hidden_states
        __a					=  output.pixel_decoder_hidden_states
        __a					=  output.transformer_decoder_hidden_states

        self.parent.assertTrue(len(__A    )  ,       len(config.backbone_config.depths    )    )
        self.parent.assertTrue(len(__A    )  ,       len(config.backbone_config.depths    )    )
        self.parent.assertTrue(len(__A    )  ,       config.decoder_layers    )


     def 	snake_case_  (	self  ,       __A  ,       __A  ,       __A  ,       __A=False    ):
        with torch.no_grad():
           __a					=  MaskaFormerModel(config=__A    )
           model.to(__A    )
           model.eval()

           __a					=  model(pixel_values=__A  ,       pixel_mask=__A    )
           __a					=  model(__A  ,       output_hidden_states=__A    )

        self.parent.assertEqual(
            output.transformer_decoder_last_hidden_state.shape  ,       (self.batch_size, self.num_queries, self.hidden_dim)  ,       )
        # let's ensure the other two hidden state exists
        self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None    )
        self.parent.assertTrue(output.encoder_last_hidden_state is not None    )

        if output_hidden_states:
           self.check_output_hidden_state(__A  ,       __A    )


     def 	snake_case_  (	self  ,       __A  ,       __A  ,       __A  ,       __A  ,       __A    ):
        __a					=  MaskaFormerForUniversalSegmentation(config=__A    )
        model.to(__A    )
        model.eval()

        def comm_check_on_output(__A    ):
           # let's still check that all the required stuff is there
           self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None    )
           self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None    )
           self.parent.assertTrue(result.encoder_last_hidden_state is not None    )
           # okay, now we need to check the logits shape
           # due to the encoder compression, masks have a //4 spatial size
           self.parent.assertEqual(
               result.masks_queries_logits.shape  ,       (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)  ,       )
           # + 1 for null class
           self.parent.assertEqual(
               result.class_queries_logits.shape  ,       (self.batch_size, self.num_queries, self.num_labels + 1)    )

        with torch.no_grad():
           __a					=  model(pixel_values=__A  ,       pixel_mask=__A    )
           __a					=  model(__A    )

           comm_check_on_output(__A    )

           __a					=  model(
               pixel_values=__A  ,       pixel_mask=__A  ,       mask_labels=__A  ,       class_labels=__A    )

        comm_check_on_output(__A    )

        self.parent.assertTrue(result.loss is not None    )
        self.parent.assertEqual(result.loss.shape  ,       torch.Size([1]    )    )







@require_torch
class     __UpperCAmelCase							(       __A ,					__A ,					unittest.TestCase							):


     """simple docstring"""
     _lowerCamelCase   =      (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
     _lowerCamelCase   =      {"""feature-extraction""": MaskaFormerModel} if is_torch_available() else {}

     _lowerCamelCase   =      False
     _lowerCamelCase   =      False
     _lowerCamelCase   =      False
     _lowerCamelCase   =      False


     def 	snake_case_  (	self    ):
        __a					=  MaskaFormerModelTester(self    )
        __a					=  ConfigTester(self  ,       config_class=__A  ,       has_text_modality=__A    )


     def 	snake_case_  (	self    ):
        self.config_tester.run_common_tests()


     def 	snake_case_  (	self    ):
        __a	,    __a					=  self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.create_and_check_maskaformer_model(__A  ,       **__A  ,       output_hidden_states=__A    )


     def 	snake_case_  (	self    ):
        __a					=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*__A    )


     @unittest.skip(reason="""Mask2Former does not use inputs_embeds"""    )
     def 	snake_case_  (	self    ):
        pass


     @unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method"""    )
     def 	snake_case_  (	self    ):
        pass


     @unittest.skip(reason="""Mask2Former is not a generative model"""    )
     def 	snake_case_  (	self    ):
        pass


     @unittest.skip(reason="""Mask2Former does not use token embeddings"""    )
     def 	snake_case_  (	self    ):
        pass


     @require_torch_multi_gpu
     @unittest.skip(
         reason="""Mask2Former has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""    )
     def 	snake_case_  (	self    ):
        pass


     @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""    )
     def 	snake_case_  (	self    ):
        pass


     def 	snake_case_  (	self    ):
        __a	,    __a					=  self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           __a					=  model_class(__A    )
           __a					=  inspect.signature(model.forward    )
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           __a					=  [*signature.parameters.keys()]

           __a					=  ["""pixel_values"""]
           self.assertListEqual(arg_names[:1]  ,       __A    )


     @slow
     def 	snake_case_  (	self    ):
        for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
           __a					=  MaskaFormerModel.from_pretrained(__A    )
           self.assertIsNotNone(__A    )


     def 	snake_case_  (	self    ):
        __a					=  (self.model_tester.min_size,) * 2
        __a					=  {
            """pixel_values""": torch.randn((2, 3, *size)  ,       device=__A    ),
            """mask_labels""": torch.randn((2, 10, *size)  ,       device=__A    ),
            """class_labels""": torch.zeros(2  ,       10  ,       device=__A    ).long(),
        }
        __a					=  self.model_tester.get_config()

        __a					=  MaskaFormerForUniversalSegmentation(__A    ).to(__A    )
        __a					=  model(**__A    )
        self.assertTrue(outputs.loss is not None    )


     def 	snake_case_  (	self    ):
        __a	,    __a					=  self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.create_and_check_maskaformer_model(__A  ,       **__A  ,       output_hidden_states=__A    )


     def 	snake_case_  (	self    ):
        __a	,    __a					=  self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           __a					=  model_class(__A    ).to(__A    )
           __a					=  model(**__A  ,       output_attentions=__A    )
           self.assertTrue(outputs.attentions is not None    )


     def 	snake_case_  (	self    ):
        if not self.model_tester.is_training:
           return

        __a					=  self.all_model_classes[1]
        __a	,    __a	,    __a	,    __a	,    __a					=  self.model_tester.prepare_config_and_inputs()

        __a					=  model_class(__A    )
        model.to(__A    )
        model.train()

        __a					=  model(__A  ,       mask_labels=__A  ,       class_labels=__A    ).loss
        loss.backward()


     def 	snake_case_  (	self    ):
        __a					=  self.all_model_classes[1]
        __a	,    __a	,    __a	,    __a	,    __a					=  self.model_tester.prepare_config_and_inputs()
        __a					=  True
        __a					=  True

        __a					=  model_class(__A    ).to(__A    )
        model.train()

        __a					=  model(__A  ,       mask_labels=__A  ,       class_labels=__A    )

        __a					=  outputs.encoder_hidden_states[0]
        encoder_hidden_states.retain_grad()

        __a					=  outputs.pixel_decoder_hidden_states[0]
        pixel_decoder_hidden_states.retain_grad()

        __a					=  outputs.transformer_decoder_hidden_states[0]
        transformer_decoder_hidden_states.retain_grad()

        __a					=  outputs.attentions[0]
        attentions.retain_grad()

        outputs.loss.backward(retain_graph=__A    )

        self.assertIsNotNone(encoder_hidden_states.grad    )
        self.assertIsNotNone(pixel_decoder_hidden_states.grad    )
        self.assertIsNotNone(transformer_decoder_hidden_states.grad    )
        self.assertIsNotNone(attentions.grad    )







SCREAMING_SNAKE_CASE							    =							1E-4




def  a   ():
   __a					=  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
   return image


@require_vision
@slow
class     __UpperCAmelCase							(       unittest.TestCase							):


     """simple docstring"""


     @cached_property
     def 	snake_case_  (	self    ):
        return "facebook/mask2former-swin-small-coco-instance"


     @cached_property
     def 	snake_case_  (	self    ):
        return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints    ) if is_vision_available() else None


     def 	snake_case_  (	self    ):
        __a					=  MaskaFormerModel.from_pretrained(self.model_checkpoints    ).to(__A    )
        __a					=  self.default_image_processor
        __a					=  prepare_img()
        __a					=  image_processor(__A  ,       return_tensors="""pt"""    ).to(__A    )
        __a					=  inputs["""pixel_values"""].shape
        # check size is divisible by 32
        self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0    )
        # check size
        self.assertEqual(__A  ,       (1, 3, 384, 384)    )

        with torch.no_grad():
           __a					=  model(**__A    )

        __a					=  torch.tensor(
            [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]    ).to(__A    )
        self.assertTrue(
            torch.allclose(
                outputs.encoder_last_hidden_state[0, 0, :3, :3]  ,       __A  ,       atol=__A    )    )

        __a					=  torch.tensor(
            [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]    ).to(__A    )
        self.assertTrue(
            torch.allclose(
                outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]  ,       __A  ,       atol=__A    )    )

        __a					=  torch.tensor(
            [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]    ).to(__A    )
        self.assertTrue(
            torch.allclose(
                outputs.transformer_decoder_last_hidden_state[0, :3, :3]  ,       __A  ,       atol=__A    )    )


     def 	snake_case_  (	self    ):
        __a					=  MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints    ).to(__A    ).eval()
        __a					=  self.default_image_processor
        __a					=  prepare_img()
        __a					=  image_processor(__A  ,       return_tensors="""pt"""    ).to(__A    )
        __a					=  inputs["""pixel_values"""].shape
        # check size is divisible by 32
        self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0    )
        # check size
        self.assertEqual(__A  ,       (1, 3, 384, 384)    )

        with torch.no_grad():
           __a					=  model(**__A    )
        # masks_queries_logits
        __a					=  outputs.masks_queries_logits
        self.assertEqual(
            masks_queries_logits.shape  ,       (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)    )
        __a					=  [
            [-8.7839, -9.0056, -8.8121],
            [-7.4104, -7.0313, -6.5401],
            [-6.6105, -6.3427, -6.4675],
        ]
        __a					=  torch.tensor(__A    ).to(__A    )
        self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]  ,       __A  ,       atol=__A    )    )
        # class_queries_logits
        __a					=  outputs.class_queries_logits
        self.assertEqual(class_queries_logits.shape  ,       (1, model.config.num_queries, model.config.num_labels + 1)    )
        __a					=  torch.tensor(
            [
                [1.8324, -8.0835, -4.1922],
                [0.8450, -9.0050, -3.6053],
                [0.3045, -7.7293, -3.0275],
            ]    ).to(__A    )
        self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]  ,       __A  ,       atol=__A    )    )


     def 	snake_case_  (	self    ):
        __a					=  MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints    ).to(__A    ).eval()
        __a					=  self.default_image_processor

        __a					=  image_processor(
            [np.zeros((3, 800, 1333)    ), np.zeros((3, 800, 1333)    )]  ,       segmentation_maps=[np.zeros((384, 384)    ).astype(np.floataa    ), np.zeros((384, 384)    ).astype(np.floataa    )]  ,       return_tensors="""pt"""  ,       )

        __a					=  inputs["""pixel_values"""].to(__A    )
        __a					=  [el.to(__A    ) for el in inputs["""mask_labels"""]]
        __a					=  [el.to(__A    ) for el in inputs["""class_labels"""]]

        with torch.no_grad():
           __a					=  model(**__A    )

        self.assertTrue(outputs.loss is not None    )










"""simple docstring"""
import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pytest

from datasets.arrow_dataset import Dataset
from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex

from .utils import require_elasticsearch, require_faiss


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







@require_faiss
class       a							(   a__  ):



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  import faiss

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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  import faiss

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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  import faiss

  lowerCAmelCase    				= self._create_dummy_dataset()
  dset.add_faiss_index_from_external_arrays(
      external_arrays=np.ones((30, 5)     ) * np.arange(30     ).reshape(-1				,       1     )				,       index_name='vecs'				,       metric_type=faiss.METRIC_INNER_PRODUCT				,       )

  # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
  # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
  # see https://bugs.python.org/issue14243 and
  # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
  with tempfile.NamedTemporaryFile(delete=_snake_case     ) as tmp_file:
   dset.save_faiss_index('vecs'				,       tmp_file.name     )
   dset.load_faiss_index('vecs2'				,       tmp_file.name     )
  os.unlink(tmp_file.name     )

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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self._create_dummy_dataset()
  dset.add_faiss_index_from_external_arrays(
      external_arrays=np.ones((30, 5)     ) * np.arange(30     ).reshape(-1				,       1     )				,       index_name='vecs'     )
  dset.drop_index('vecs'     )
  self.assertRaises(_snake_case				,       partial(dset.get_nearest_examples				,       'vecs2'				,       np.ones(5				,       dtype=np.floataa     )     )     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  from elasticsearch import Elasticsearch

  lowerCAmelCase    				= self._create_dummy_dataset()
  with patch('elasticsearch.Elasticsearch.search'     ) as mocked_search, patch(
      'elasticsearch.client.IndicesClient.create'     ) as mocked_index_create, patch('elasticsearch.helpers.streaming_bulk'     ) as mocked_bulk:
   lowerCAmelCase    				= {'acknowledged': True}
   mocked_bulk.return_value([(True, None)] * 30     )
   lowerCAmelCase    				= {'hits': {'hits': [{'_score': 1, '_id': 29}]}}
   lowerCAmelCase    				= Elasticsearch()

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







@require_faiss
class       a							(   a__  ):



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  import faiss

  lowerCAmelCase    				= FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT     )

  # add vectors
  index.add_vectors(np.eye(5				,       dtype=np.floataa     )     )
  self.assertIsNotNone(index.faiss_index     )
  self.assertEqual(index.faiss_index.ntotal				,       5     )
  index.add_vectors(np.zeros((5, 5)				,       dtype=np.floataa     )     )
  self.assertEqual(index.faiss_index.ntotal				,       10     )

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

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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  import faiss

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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  import faiss

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



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  import faiss

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

  # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
  # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
  # see https://bugs.python.org/issue14243 and
  # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
  with tempfile.NamedTemporaryFile(delete=_snake_case     ) as tmp_file:
   index.save(tmp_file.name     )
   lowerCAmelCase    				= FaissIndex.load(tmp_file.name     )
  os.unlink(tmp_file.name     )

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


@require_faiss
def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Dict			):
 import faiss

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

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

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







@require_elasticsearch
class       a							(   a__  ):



 def    UpperCamelCase__  (     self     ):







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

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

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

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

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





