



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


_A				:    Optional[Any]     						=pytest.mark.integration



@require_faiss
class 		lowerCamelCase__      (						A     ):




					'''simple docstring'''



					def 						__UpperCAmelCase     (  self				:	int					)  ->       str:



						'''simple docstring'''





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




					def 						__UpperCAmelCase     (  self				:	Optional[int]					)  ->       int:



						'''simple docstring'''





						import faiss

						_lowercase					:       Dataset       =		self._create_dummy_dataset()
						_lowercase					:       Union[str, Any]       =		dset.map(
						    lambda UpperCamelCase_     ,  UpperCamelCase_					: {"vecs": i * np.ones(5     ,  dtype=np.floataa					)}     ,  with_indices=UpperCamelCase_     ,  keep_in_memory=UpperCamelCase_					)
						_lowercase					:       List[str]       =		dset.add_faiss_index('vecs'     ,  batch_size=100     ,  metric_type=faiss.METRIC_INNER_PRODUCT					)
						_lowercase  ,    _lowercase					:       Any       =		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				:	Any					)  ->       int:



						'''simple docstring'''





						import faiss

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




					def 						__UpperCAmelCase     (  self				:	str					)  ->       List[str]:



						'''simple docstring'''





						import faiss

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

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




					def 						__UpperCAmelCase     (  self				:	Any					)  ->       Union[str, Any]:



						'''simple docstring'''





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




					def 						__UpperCAmelCase     (  self				:	List[Any]					)  ->       Dict:



						'''simple docstring'''





						from elasticsearch import Elasticsearch

						_lowercase					:       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:
							_lowercase					:       Optional[int]       =		{'acknowledged': True}
							mocked_bulk.return_value([(True, None)] * 30					)
							_lowercase					:       Union[str, Any]       =		{'hits': {'hits': [{'_score': 1, '_id': 29}]}}
							_lowercase					:       str       =		Elasticsearch()

							dset.add_elasticsearch_index('filename'     ,  es_client=UpperCamelCase_					)
							_lowercase  ,    _lowercase					:       int       =		dset.get_nearest_examples('filename'     ,  'my_name-train_29'					)
							self.assertEqual(examples['filename'][0]     ,  'my_name-train_29'					)






@require_faiss
class 		lowerCamelCase__      (						A     ):




					'''simple docstring'''



					def 						__UpperCAmelCase     (  self				:	Tuple					)  ->       str:



						'''simple docstring'''





						import faiss

						_lowercase					:       List[str]       =		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
						_lowercase					:       List[str]       =		np.zeros(5     ,  dtype=np.floataa					)
						_lowercase					:       int       =		1
						_lowercase  ,    _lowercase					:       Optional[Any]       =		index.search(UpperCamelCase_					)
						self.assertRaises(UpperCamelCase_     ,  index.search     ,  query.reshape(-1     ,  1					)					)
						self.assertGreater(scores[0]     ,  0					)
						self.assertEqual(indices[0]     ,  1					)

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




					def 						__UpperCAmelCase     (  self				:	List[str]					)  ->       str:



						'''simple docstring'''





						import faiss

						_lowercase					:       List[Any]       =		FaissIndex(string_factory='Flat'					)
						index.add_vectors(np.eye(5     ,  dtype=np.floataa					)					)
						self.assertIsInstance(index.faiss_index     ,  faiss.IndexFlat					)
						_lowercase					:       int       =		FaissIndex(string_factory='LSH'					)
						index.add_vectors(np.eye(5     ,  dtype=np.floataa					)					)
						self.assertIsInstance(index.faiss_index     ,  faiss.IndexLSH					)
						with self.assertRaises(UpperCamelCase_					):
							_lowercase					:       Union[str, Any]       =		FaissIndex(string_factory='Flat'     ,  custom_index=faiss.IndexFlat(5					)					)




					def 						__UpperCAmelCase     (  self				:	int					)  ->       int:



						'''simple docstring'''





						import faiss

						_lowercase					:       Tuple       =		faiss.IndexFlat(5					)
						_lowercase					:       List[Any]       =		FaissIndex(custom_index=UpperCamelCase_					)
						index.add_vectors(np.eye(5     ,  dtype=np.floataa					)					)
						self.assertIsInstance(index.faiss_index     ,  faiss.IndexFlat					)




					def 						__UpperCAmelCase     (  self				:	str					)  ->       Any:



						'''simple docstring'''





						import faiss

						_lowercase					:       List[Any]       =		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=UpperCamelCase_					) as tmp_file:
							index.save(tmp_file.name					)
							_lowercase					:       List[Any]       =		FaissIndex.load(tmp_file.name					)
						os.unlink(tmp_file.name					)

						_lowercase					:       Dict       =		np.zeros(5     ,  dtype=np.floataa					)
						_lowercase					:       List[str]       =		1
						_lowercase  ,    _lowercase					:       Dict       =		index.search(UpperCamelCase_					)
						self.assertGreater(scores[0]     ,  0					)
						self.assertEqual(indices[0]     ,  1					)







@require_faiss
def 						__UpperCamelCase				(							_lowercase					) ->			Optional[int]:
	import faiss

	_lowercase					:       Optional[int]       =		FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT					)
	index.add_vectors(np.eye(5,				dtype=np.floataa					)					)

	_lowercase					:       str       =		'index.faiss'
	_lowercase					:       List[Any]       =		f'''mock://{index_name}'''
	index.save(_lowercase,				storage_options=mockfs.storage_options					)
	_lowercase					:       Any       =		FaissIndex.load(_lowercase,				storage_options=mockfs.storage_options					)

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






@require_elasticsearch
class 		lowerCamelCase__      (						A     ):




					'''simple docstring'''



					def 						__UpperCAmelCase     (  self				:	List[str]					)  ->       List[Any]:



						'''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:
							_lowercase					:       Tuple       =		Elasticsearch()
							_lowercase					:       Optional[int]       =		{'acknowledged': True}
							_lowercase					:       List[Any]       =		ElasticSearchIndex(es_client=UpperCamelCase_					)
							mocked_bulk.return_value([(True, None)] * 3					)
							index.add_documents(['foo', 'bar', 'foobar']					)

							# single query
							_lowercase					:       Optional[Any]       =		'foo'
							_lowercase					:       Optional[Any]       =		{'hits': {'hits': [{'_score': 1, '_id': 0}]}}
							_lowercase  ,    _lowercase					:       Dict       =		index.search(UpperCamelCase_					)
							self.assertEqual(scores[0]     ,  1					)
							self.assertEqual(indices[0]     ,  0					)

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

							# batched queries
							_lowercase					:       Dict       =		['foo', 'bar', 'foobar']
							_lowercase					:       List[Any]       =		{'hits': {'hits': [{'_score': 1, '_id': 1}]}}
							_lowercase  ,    _lowercase					:       Tuple       =		index.search_batch(UpperCamelCase_					)
							_lowercase					:       Union[str, Any]       =		[scores[0] for scores in total_scores]
							_lowercase					:       str       =		[indices[0] for indices in total_indices]
							self.assertGreater(np.min(UpperCamelCase_					)     ,  0					)
							self.assertListEqual([1, 1, 1]     ,  UpperCamelCase_					)

							# batched queries with timeout
							_lowercase					:       List[str]       =		['foo', 'bar', 'foobar']
							_lowercase					:       Optional[int]       =		{'hits': {'hits': [{'_score': 1, '_id': 1}]}}
							_lowercase  ,    _lowercase					:       Tuple       =		index.search_batch(UpperCamelCase_     ,  request_timeout=30					)
							_lowercase					:       Optional[Any]       =		[scores[0] for scores in total_scores]
							_lowercase					:       Any       =		[indices[0] for indices in total_indices]
							self.assertGreater(np.min(UpperCamelCase_					)     ,  0					)
							self.assertListEqual([1, 1, 1]     ,  UpperCamelCase_					)





'''simple docstring'''

import unittest

from transformers import AutoTokenizer, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow


if is_flax_available():
		import jax.numpy as jnp

		from transformers import FlaxXLMRobertaModel



@require_sentencepiece
@require_tokenizers
@require_flax
class 		lowerCamelCase__      (						unittest.TestCase     ):




					'''simple docstring'''



					@slow
					def 						__UpperCAmelCase     (  self				:	int					)  ->       Union[str, Any]:



						'''simple docstring'''





						_lowercase					:       List[Any]       =		FlaxXLMRobertaModel.from_pretrained('xlm-roberta-base'					)
						_lowercase					:       str       =		AutoTokenizer.from_pretrained('xlm-roberta-base'					)
						_lowercase					:       List[Any]       =		'The dog is cute and lives in the garden house'
						_lowercase					:       Optional[int]       =		jnp.array([tokenizer.encode(UpperCamelCase_					)]					)

						_lowercase					:       int       =		(1, 12, 768)  # batch_size, sequence_length, embedding_vector_dim
						_lowercase					:       Tuple       =		jnp.array(
						    [[-0.01_01, 0.12_18, -0.08_03, 0.08_01, 0.13_27, 0.07_76, -0.12_15, 0.23_83, 0.33_38, 0.31_06, 0.03_00, 0.02_52]]					)

						_lowercase					:       List[str]       =		model(UpperCamelCase_					)['last_hidden_state']
						self.assertEqual(output.shape     ,  UpperCamelCase_					)
						# compare the actual values for a slice of last dim
						self.assertTrue(jnp.allclose(output[:, :, -1]     ,  UpperCamelCase_     ,  atol=1E-3					)					)
