

import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import (
       FocalNetBackbone,
       FocalNetForImageClassification,
       FocalNetForMaskedImageModeling,
       FocalNetModel,
   )
   from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor





class        lowerCAmelCase__  :



       '''simple docstring'''







       def __init__(							self ,   lowercase__ ,   lowercase__=1_3 ,   lowercase__=3_2 ,   lowercase__=2 ,   lowercase__=3 ,   lowercase__=1_6 ,   lowercase__=[3_2, 6_4, 1_2_8] ,   lowercase__=[1, 2, 1] ,   lowercase__=[2, 2, 4] ,   lowercase__=2 ,   lowercase__=2.0 ,   lowercase__=True ,   lowercase__=0.0 ,   lowercase__=0.0 ,   lowercase__=0.1 ,   lowercase__="gelu" ,   lowercase__=False ,   lowercase__=True ,   lowercase__=0.02 ,   lowercase__=1E-5 ,   lowercase__=True ,   lowercase__=None ,   lowercase__=True ,   lowercase__=1_0 ,   lowercase__=8 ,   lowercase__=["stage1", "stage2"] ,   lowercase__=[1, 2] ,   ):







             '''simple docstring'''



             __A    =parent
             __A    =batch_size
             __A    =image_size
             __A    =patch_size
             __A    =num_channels
             __A    =embed_dim
             __A    =hidden_sizes
             __A    =depths
             __A    =num_heads
             __A    =window_size
             __A    =mlp_ratio
             __A    =qkv_bias
             __A    =hidden_dropout_prob
             __A    =attention_probs_dropout_prob
             __A    =drop_path_rate
             __A    =hidden_act
             __A    =use_absolute_embeddings
             __A    =patch_norm
             __A    =layer_norm_eps
             __A    =initializer_range
             __A    =is_training
             __A    =scope
             __A    =use_labels
             __A    =type_sequence_label_size
             __A    =encoder_stride
             __A    =out_features
             __A    =out_indices

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

             __A    =None
             if self.use_labels:
                   __A    =ids_tensor([self.batch_size] ,   self.type_sequence_label_size     )

             __A    =self.get_config()

             return config, pixel_values, labels

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             return FocalNetConfig(
                 image_size=self.image_size ,   patch_size=self.patch_size ,   num_channels=self.num_channels ,   embed_dim=self.embed_dim ,   hidden_sizes=self.hidden_sizes ,   depths=self.depths ,   num_heads=self.num_heads ,   window_size=self.window_size ,   mlp_ratio=self.mlp_ratio ,   qkv_bias=self.qkv_bias ,   hidden_dropout_prob=self.hidden_dropout_prob ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob ,   drop_path_rate=self.drop_path_rate ,   hidden_act=self.hidden_act ,   use_absolute_embeddings=self.use_absolute_embeddings ,   path_norm=self.patch_norm ,   layer_norm_eps=self.layer_norm_eps ,   initializer_range=self.initializer_range ,   encoder_stride=self.encoder_stride ,   out_features=self.out_features ,   out_indices=self.out_indices ,   )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =FocalNetModel(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =model(lowercase__     )

             __A    =((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths     ) - 1))
             __A    =int(config.embed_dim * 2 ** (len(config.depths     ) - 1)     )

             self.parent.assertEqual(result.last_hidden_state.shape ,   (self.batch_size, expected_seq_len, expected_dim)     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =FocalNetBackbone(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =model(lowercase__     )

             # verify feature maps
             self.parent.assertEqual(len(result.feature_maps     ) ,   len(config.out_features     )     )
             self.parent.assertListEqual(list(result.feature_maps[0].shape     ) ,   [self.batch_size, self.image_size, 8, 8]     )

             # verify channels
             self.parent.assertEqual(len(model.channels     ) ,   len(config.out_features     )     )
             self.parent.assertListEqual(model.channels ,   config.hidden_sizes[:-1]     )

             # verify backbone works with out_features=None
             __A    =None
             __A    =FocalNetBackbone(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =model(lowercase__     )

             # verify feature maps
             self.parent.assertEqual(len(result.feature_maps     ) ,   1     )
             self.parent.assertListEqual(list(result.feature_maps[0].shape     ) ,   [self.batch_size, self.image_size * 2, 4, 4]     )

             # verify channels
             self.parent.assertEqual(len(model.channels     ) ,   1     )
             self.parent.assertListEqual(model.channels ,   [config.hidden_sizes[-1]]     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =FocalNetForMaskedImageModeling(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =model(lowercase__     )
             self.parent.assertEqual(
                 result.reconstruction.shape ,   (self.batch_size, self.num_channels, self.image_size, self.image_size)     )

             # test greyscale images
             __A    =1
             __A    =FocalNetForMaskedImageModeling(lowercase__     )
             model.to(lowercase__     )
             model.eval()

             __A    =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
             __A    =model(lowercase__     )
             self.parent.assertEqual(result.reconstruction.shape ,   (self.batch_size, 1, self.image_size, self.image_size)     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =self.type_sequence_label_size
             __A    =FocalNetForImageClassification(lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =model(lowercase__ ,   labels=lowercase__     )
             self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.type_sequence_label_size)     )

             # test greyscale images
             __A    =1
             __A    =FocalNetForImageClassification(lowercase__     )
             model.to(lowercase__     )
             model.eval()

             __A    =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
             __A    =model(lowercase__     )
             self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.type_sequence_label_size)     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.prepare_config_and_inputs()

             __A    ,		__A    ,		__A    =config_and_inputs
             __A    ={'''pixel_values''': pixel_values}
             return config, inputs_dict





@require_torch
class        lowerCAmelCase__  (					__magic_name__				,							__magic_name__				,							unittest.TestCase			):



       '''simple docstring'''







       lowercase_						=    (
           (
               FocalNetModel,
               FocalNetForImageClassification,
               FocalNetForMaskedImageModeling,
               FocalNetBackbone,
           )
           if is_torch_available()
           else ()
       )
       lowercase_						=    (
           {"""feature-extraction""": FocalNetModel, """image-classification""": FocalNetForImageClassification}
           if is_torch_available()
           else {}
       )
       lowercase_						=    False

       lowercase_						=    False
       lowercase_						=    False
       lowercase_						=    False
       lowercase_						=    False

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =FocalNetModelTester(self     )
             __A    =ConfigTester(self ,   config_class=lowercase__ ,   embed_dim=3_7 ,   has_text_modality=lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             self.create_and_test_config_common_properties()
             self.config_tester.create_and_test_config_to_json_string()
             self.config_tester.create_and_test_config_to_json_file()
             self.config_tester.create_and_test_config_from_and_save_pretrained()
             self.config_tester.create_and_test_config_with_num_labels()
             self.config_tester.check_config_can_be_init_without_params()
             self.config_tester.check_config_arguments_init()

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             return

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_backbone(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_image_modeling(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*lowercase__     )

       @unittest.skip(reason='''FocalNet does not use inputs_embeds'''     )
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             pass

       @unittest.skip(reason='''FocalNet does not use feedforward chunking'''     )
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             pass

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ,		__A    =self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes[:-1]:
                   __A    =model_class(lowercase__     )
                   self.assertIsInstance(model.get_input_embeddings() ,   (nn.Module)     )
                   __A    =model.get_output_embeddings()
                   self.assertTrue(x is None or isinstance(lowercase__ ,   nn.Linear     )     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ,		__A    =self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes[:-1]:
                   __A    =model_class(lowercase__     )
                   __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] ,   lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =model_class(lowercase__     )
             model.to(lowercase__     )
             model.eval()

             with torch.no_grad():
                   __A    =model(**self._prepare_for_class(lowercase__ ,   lowercase__     )     )

             __A    =outputs.hidden_states

             __A    =getattr(
                 self.model_tester ,   '''expected_num_hidden_layers''' ,   len(self.model_tester.depths     ) + 1     )
             self.assertEqual(len(lowercase__     ) ,   lowercase__     )

             # FocalNet has a different seq_length
             __A    =(
                 config.patch_size
                 if isinstance(config.patch_size ,   collections.abc.Iterable     )
                 else (config.patch_size, config.patch_size)
             )

             __A    =(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

             self.assertListEqual(
                 list(hidden_states[0].shape[-2:]     ) ,   [num_patches, self.model_tester.embed_dim] ,   )

             __A    =outputs.reshaped_hidden_states
             self.assertEqual(len(lowercase__     ) ,   lowercase__     )

             __A    ,		__A    ,		__A    ,		__A    =reshaped_hidden_states[0].shape
             __A    =(
                 reshaped_hidden_states[0].view(lowercase__ ,   lowercase__ ,   height * width     ).permute(0 ,   2 ,   1     )
             )
             self.assertListEqual(
                 list(reshaped_hidden_states.shape[-2:]     ) ,   [num_patches, self.model_tester.embed_dim] ,   )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ,		__A    =self.model_tester.prepare_config_and_inputs_for_common()

             __A    =(
                 self.model_tester.image_size
                 if isinstance(self.model_tester.image_size ,   collections.abc.Iterable     )
                 else (self.model_tester.image_size, self.model_tester.image_size)
             )

             for model_class in self.all_model_classes[:-1]:
                   __A    =True
                   self.check_hidden_states_output(lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     )

                   # check that output_hidden_states also work using config
                   del inputs_dict["output_hidden_states"]
                   __A    =True

                   self.check_hidden_states_output(lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ,		__A    =self.model_tester.prepare_config_and_inputs_for_common()
             __A    =3

             __A    =(
                 self.model_tester.image_size
                 if isinstance(self.model_tester.image_size ,   collections.abc.Iterable     )
                 else (self.model_tester.image_size, self.model_tester.image_size)
             )
             __A    =(
                 config.patch_size
                 if isinstance(config.patch_size ,   collections.abc.Iterable     )
                 else (config.patch_size, config.patch_size)
             )

             __A    =image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
             __A    =image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

             for model_class in self.all_model_classes[:-1]:
                   __A    =True
                   self.check_hidden_states_output(lowercase__ ,   lowercase__ ,   lowercase__ ,   (padded_height, padded_width)     )

                   # check that output_hidden_states also work using config
                   del inputs_dict["output_hidden_states"]
                   __A    =True
                   self.check_hidden_states_output(lowercase__ ,   lowercase__ ,   lowercase__ ,   (padded_height, padded_width)     )

       @slow
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   __A    =FocalNetModel.from_pretrained(lowercase__     )
                   self.assertIsNotNone(lowercase__     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ,		__A    =self.model_tester.prepare_config_and_inputs_for_common()

             __A    =_config_zero_init(lowercase__     )
             for model_class in self.all_model_classes:
                   __A    =model_class(config=lowercase__     )
                   for name, param in model.named_parameters():
                         if "embeddings" not in name and param.requires_grad:
                               self.assertIn(
                                   ((param.data.mean() * 1E9).round() / 1E9).item() ,   [0.0, 1.0] ,   msg=f'''Parameter {name} of model {model_class} seems not properly initialized''' ,   )





@require_vision
@require_torch
class        lowerCAmelCase__  (					unittest.TestCase			):



       '''simple docstring'''







       @cached_property
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             return AutoImageProcessor.from_pretrained('''microsoft/focalnet-tiny'''     ) if is_vision_available() else None


       @slow
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =FocalNetForImageClassification.from_pretrained('''microsoft/focalnet-tiny'''     ).to(lowercase__     )
             __A    =self.default_image_processor

             __A    =Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''     )
             __A    =image_processor(images=lowercase__ ,   return_tensors='''pt'''     ).to(lowercase__     )

             # forward pass
             with torch.no_grad():
                   __A    =model(**lowercase__     )

             # verify the logits
             __A    =torch.Size((1, 1_0_0_0)     )
             self.assertEqual(outputs.logits.shape ,   lowercase__     )
             __A    =torch.tensor([0.2166, -0.4368, 0.2191]     ).to(lowercase__     )
             self.assertTrue(torch.allclose(outputs.logits[0, :3] ,   lowercase__ ,   atol=1E-4     )     )
             self.assertTrue(outputs.logits.argmax(dim=-1     ).item() ,   2_8_1     )





@require_torch
class        lowerCAmelCase__  (					__magic_name__				,							unittest.TestCase			):



       '''simple docstring'''







       lowercase_						=    (FocalNetBackbone,) if is_torch_available() else ()
       lowercase_						=    FocalNetConfig

       lowercase_						=    False

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =FocalNetModelTester(self     )





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


_lowerCamelCase      :					List[str]					    =      pytest.mark.integration





@require_faiss
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

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

             __A    =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 ,   )
             __A    ,		__A    =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

             __A    =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=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     )

             __A    ,		__A    =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'''



             __A    =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(lowercase__ ,   partial(dset.get_nearest_examples ,   '''vecs2''' ,   np.ones(5 ,   dtype=np.floataa     )     )     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             from elasticsearch import Elasticsearch

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

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





@require_faiss
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

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

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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

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


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =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     )
                   __A    =FaissIndex.load(tmp_file.name     )
             os.unlink(tmp_file.name     )

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





@require_faiss
def      A__					(			__A  :						Dict			)					->List[Any]:
      import faiss

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

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

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





@require_elasticsearch
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







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

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

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

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

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


