


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


snake_case_    					=	pytest.mark.integration







@require_faiss
class 				A_      (					SCREAMING_SNAKE_CASE_    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:Dict )    ->      List[Any]:
												UpperCAmelCase					=			Dataset.from_dict({'filename': ['my_name-train' + '_' + str(lowercase_ ) for x in np.arange(30 ).tolist()]} )
												return dset






						def      UpperCAmelCase__					( self			:int )    ->      Tuple:
												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				,   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      UpperCAmelCase__					( self			:Dict )    ->      Union[str, Any]:
												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				,   UpperCAmelCase					=			dset.get_nearest_examples('vecs'       ,  np.ones(5       ,  dtype=np.floataa ) )
												self.assertEqual(examples['filename'][0]       ,  'my_name-train_29' )






						def      UpperCAmelCase__					( self			:Tuple )    ->      Any:
												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				,   UpperCAmelCase					=			dset.get_nearest_examples('vecs2'       ,  np.ones(5       ,  dtype=np.floataa ) )
												self.assertEqual(examples['filename'][0]       ,  'my_name-train_29' )






						def      UpperCAmelCase__					( self			:Optional[Any] )    ->      List[Any]:
												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      UpperCAmelCase__					( self			:List[str] )    ->      Optional[Any]:
												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				,   UpperCAmelCase					=			dset.get_nearest_examples('filename'       ,  'my_name-train_29' )
																		self.assertEqual(examples['filename'][0]       ,  'my_name-train_29' )







@require_faiss
class 				A_      (					SCREAMING_SNAKE_CASE_    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:Union[str, Any] )    ->      List[Any]:
												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				,   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				,   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      UpperCAmelCase__					( self			:int )    ->      List[Any]:
												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      UpperCAmelCase__					( self			:str )    ->      str:
												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      UpperCAmelCase__					( self			:Optional[int] )    ->      Optional[int]:
												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				,   UpperCAmelCase					=			index.search(lowercase_ )
												self.assertGreater(scores[0]       ,  0 )
												self.assertEqual(indices[0]       ,  1 )


@require_faiss
def 			_lowerCAmelCase				(       lowercase_						):
						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(lowercase_		,     storage_options=mockfs.storage_options						)
						UpperCAmelCase					=			FaissIndex.load(lowercase_		,     storage_options=mockfs.storage_options						)

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







@require_elasticsearch
class 				A_      (					SCREAMING_SNAKE_CASE_    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:List[str] )    ->      List[str]:
												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				,   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				,   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				,   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				,   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_ )










"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


snake_case_    					=	{
    """configuration_rembert""": ["""REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """RemBertConfig""", """RemBertOnnxConfig"""]
}

try:
				if not is_sentencepiece_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				snake_case_    					=	["""RemBertTokenizer"""]

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				snake_case_    					=	["""RemBertTokenizerFast"""]

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				snake_case_    					=	[
				    """REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
				    """RemBertForCausalLM""",
				    """RemBertForMaskedLM""",
				    """RemBertForMultipleChoice""",
				    """RemBertForQuestionAnswering""",
				    """RemBertForSequenceClassification""",
				    """RemBertForTokenClassification""",
				    """RemBertLayer""",
				    """RemBertModel""",
				    """RemBertPreTrainedModel""",
				    """load_tf_weights_in_rembert""",
				]


try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				snake_case_    					=	[
				    """TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
				    """TFRemBertForCausalLM""",
				    """TFRemBertForMaskedLM""",
				    """TFRemBertForMultipleChoice""",
				    """TFRemBertForQuestionAnswering""",
				    """TFRemBertForSequenceClassification""",
				    """TFRemBertForTokenClassification""",
				    """TFRemBertLayer""",
				    """TFRemBertModel""",
				    """TFRemBertPreTrainedModel""",
				]


if TYPE_CHECKING:
				from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig

				try:
								if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_rembert import RemBertTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_rembert_fast import RemBertTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_rembert import (
								    REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    RemBertForCausalLM,
								    RemBertForMaskedLM,
								    RemBertForMultipleChoice,
								    RemBertForQuestionAnswering,
								    RemBertForSequenceClassification,
								    RemBertForTokenClassification,
								    RemBertLayer,
								    RemBertModel,
								    RemBertPreTrainedModel,
								    load_tf_weights_in_rembert,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_rembert import (
								    TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFRemBertForCausalLM,
								    TFRemBertForMaskedLM,
								    TFRemBertForMultipleChoice,
								    TFRemBertForQuestionAnswering,
								    TFRemBertForSequenceClassification,
								    TFRemBertForTokenClassification,
								    TFRemBertLayer,
								    TFRemBertModel,
								    TFRemBertPreTrainedModel,
								)


else:
				import sys

				snake_case_    					=	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





