


"""simple docstring"""







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowerCAmelCase__      				=   {
    '''configuration_groupvit''': [
        '''GROUPVIT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''GroupViTConfig''',
        '''GroupViTOnnxConfig''',
        '''GroupViTTextConfig''',
        '''GroupViTVisionConfig''',
    ],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCAmelCase__      				=   [
         '''GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''GroupViTModel''',
         '''GroupViTPreTrainedModel''',
         '''GroupViTTextModel''',
         '''GroupViTVisionModel''',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCAmelCase__      				=   [
         '''TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TFGroupViTModel''',
         '''TFGroupViTPreTrainedModel''',
         '''TFGroupViTTextModel''',
         '''TFGroupViTVisionModel''',
     ]

if TYPE_CHECKING:
     from .configuration_groupvit import (
         GROUPVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
         GroupViTConfig,
         GroupViTOnnxConfig,
         GroupViTTextConfig,
         GroupViTVisionConfig,
     )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_groupvit import (
              GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
              GroupViTModel,
              GroupViTPreTrainedModel,
              GroupViTTextModel,
              GroupViTVisionModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_groupvit import (
              TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFGroupViTModel,
              TFGroupViTPreTrainedModel,
              TFGroupViTTextModel,
              TFGroupViTVisionModel,
          )

else:
     import sys

     lowerCAmelCase__      				=   _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)









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


lowerCAmelCase__      				=   pytest.mark.integration


@require_faiss
class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       def        lowercase__							(						self						):






         """simple docstring"""





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

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Dataset								=							self._create_dummy_dataset()
         lowerCAmelCase     :     Union[str, Any]								=							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     :     Union[str, Any]								=							dset.add_faiss_index("vecs"		,     batch_size=100		,     metric_type=faiss.METRIC_INNER_PRODUCT						)
         lowerCAmelCase						,  lowerCAmelCase     :     Tuple								=							dset.get_nearest_examples("vecs"		,     np.ones(5		,     dtype=np.floataa						)						)
         self.assertEqual(examples["filename"][0]		,     "my_name-train_29"						)
         dset.drop_index("vecs"						)

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Dataset								=							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=100		,     metric_type=faiss.METRIC_INNER_PRODUCT		,     )
         lowerCAmelCase						,  lowerCAmelCase     :     Optional[Any]								=							dset.get_nearest_examples("vecs"		,     np.ones(5		,     dtype=np.floataa						)						)
         self.assertEqual(examples["filename"][0]		,     "my_name-train_29"						)

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Dataset								=							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        lowercase__							(						self						):






         """simple docstring"""





         from elasticsearch import Elasticsearch

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

           dset.add_elasticsearch_index("filename"		,     es_client=snake_case__						)
           lowerCAmelCase						,  lowerCAmelCase     :     int								=							dset.get_nearest_examples("filename"		,     "my_name-train_29"						)
           self.assertEqual(examples["filename"][0]		,     "my_name-train_29"						)


@require_faiss
class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Optional[int]								=							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     :     int								=							np.zeros(5		,     dtype=np.floataa						)
         lowerCAmelCase     :     Optional[int]								=							1
         lowerCAmelCase						,  lowerCAmelCase     :     Optional[Any]								=							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     :     Union[str, Any]								=							np.eye(5		,     dtype=np.floataa						)[::-1]
         lowerCAmelCase						,  lowerCAmelCase     :     str								=							index.search_batch(snake_case__						)
         self.assertRaises(snake_case__		,     index.search_batch		,     queries[0]						)
         lowerCAmelCase     :     Optional[int]								=							[scores[0] for scores in total_scores]
         lowerCAmelCase     :     int								=							[indices[0] for indices in total_indices]
         self.assertGreater(np.min(snake_case__						)		,     0						)
         self.assertListEqual([4, 3, 2, 1, 0]		,     snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Dict								=							FaissIndex(string_factory="Flat"						)
         index.add_vectors(np.eye(5		,     dtype=np.floataa						)						)
         self.assertIsInstance(index.faiss_index		,     faiss.IndexFlat						)
         lowerCAmelCase     :     Union[str, Any]								=							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     :     List[Any]								=							FaissIndex(string_factory="Flat"		,     custom_index=faiss.IndexFlat(5						)						)

       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     Any								=							faiss.IndexFlat(5						)
         lowerCAmelCase     :     Union[str, Any]								=							FaissIndex(custom_index=snake_case__						)
         index.add_vectors(np.eye(5		,     dtype=np.floataa						)						)
         self.assertIsInstance(index.faiss_index		,     faiss.IndexFlat						)







       def        lowercase__							(						self						):






         """simple docstring"""





         import faiss

         lowerCAmelCase     :     str								=							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     :     int								=							FaissIndex.load(tmp_file.name						)
         os.unlink(tmp_file.name						)

         lowerCAmelCase     :     Union[str, Any]								=							np.zeros(5		,     dtype=np.floataa						)
         lowerCAmelCase     :     List[str]								=							1
         lowerCAmelCase						,  lowerCAmelCase     :     Tuple								=							index.search(snake_case__						)
         self.assertGreater(scores[0]		,     0						)
         self.assertEqual(indices[0]		,     1						)

@require_faiss
def 					a__			(				SCREAMING_SNAKE_CASE    :    Any				):


  '''simple docstring'''







  import faiss

  lowerCAmelCase     :     List[Any]								=							FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT				)
  index.add_vectors(np.eye(5				,      dtype=np.floataa				)				)

  lowerCAmelCase     :     Union[str, Any]								=							"index.faiss"
  lowerCAmelCase     :     List[str]								=							f"""mock://{index_name}"""
  index.save(SCREAMING_SNAKE_CASE				,      storage_options=mockfs.storage_options				)
  lowerCAmelCase     :     Optional[Any]								=							FaissIndex.load(SCREAMING_SNAKE_CASE				,      storage_options=mockfs.storage_options				)

  lowerCAmelCase     :     Optional[int]								=							np.zeros(5				,      dtype=np.floataa				)
  lowerCAmelCase     :     Any								=							1
  lowerCAmelCase						,  lowerCAmelCase     :     Optional[int]								=							index.search(SCREAMING_SNAKE_CASE				)
  assert scores[0] > 0
  assert indices[0] == 1





@require_elasticsearch
class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       def        lowercase__							(						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     :     List[str]								=							Elasticsearch()
           lowerCAmelCase     :     Dict								=							{"acknowledged": True}
           lowerCAmelCase     :     Optional[int]								=							ElasticSearchIndex(es_client=snake_case__						)
           mocked_bulk.return_value([(True, None)] * 3						)
           index.add_documents(["foo", "bar", "foobar"]						)

           # single query
           lowerCAmelCase     :     List[str]								=							"foo"
           lowerCAmelCase     :     List[str]								=							{"hits": {"hits": [{"_score": 1, "_id": 0}]}}
           lowerCAmelCase						,  lowerCAmelCase     :     Optional[int]								=							index.search(snake_case__						)
           self.assertEqual(scores[0]		,     1						)
           self.assertEqual(indices[0]		,     0						)

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

           # batched queries
           lowerCAmelCase     :     Any								=							["foo", "bar", "foobar"]
           lowerCAmelCase     :     Optional[int]								=							{"hits": {"hits": [{"_score": 1, "_id": 1}]}}
           lowerCAmelCase						,  lowerCAmelCase     :     Any								=							index.search_batch(snake_case__						)
           lowerCAmelCase     :     Tuple								=							[scores[0] for scores in total_scores]
           lowerCAmelCase     :     Dict								=							[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     :     Optional[Any]								=							["foo", "bar", "foobar"]
           lowerCAmelCase     :     List[Any]								=							{"hits": {"hits": [{"_score": 1, "_id": 1}]}}
           lowerCAmelCase						,  lowerCAmelCase     :     Any								=							index.search_batch(snake_case__		,     request_timeout=30						)
           lowerCAmelCase     :     Dict								=							[scores[0] for scores in total_scores]
           lowerCAmelCase     :     int								=							[indices[0] for indices in total_indices]
           self.assertGreater(np.min(snake_case__						)		,     0						)
           self.assertListEqual([1, 1, 1]		,     snake_case__						)





