




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


lowercase__     						=						pytest.mark.integration






@require_faiss
class lowerCAmelCase__	(     lowercase						):






			'''simple docstring'''



			def 					A_    (  self  ):
							_lowerCamelCase     :			int    =    Dataset.from_dict({'filename': ['my_name-train' + '_' + str(lowercase  ) for x in np.arange(30  ).tolist()]}  )
							return dset
			def 					A_    (  self  ):
							import faiss

							_lowerCamelCase     :			Dataset    =    self._create_dummy_dataset()
							_lowerCamelCase     :			str    =    dset.map(
							    lambda lowercase  ,      lowercase  : {"vecs": i * np.ones(5  ,      dtype=np.floataa  )}  ,      with_indices=lowercase  ,      keep_in_memory=lowercase  )
							_lowerCamelCase     :			Optional[Any]    =    dset.add_faiss_index('vecs'  ,      batch_size=100  ,      metric_type=faiss.METRIC_INNER_PRODUCT  )
							_lowerCamelCase, _lowerCamelCase     :			int    =    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 					A_    (  self  ):
							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[int]    =    dset.get_nearest_examples('vecs'  ,      np.ones(5  ,      dtype=np.floataa  )  )
							self.assertEqual(examples['filename'][0]  ,      'my_name-train_29'  )
			def 					A_    (  self  ):
							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=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  )

							_lowerCamelCase, _lowerCamelCase     :			Dict    =    dset.get_nearest_examples('vecs2'  ,      np.ones(5  ,      dtype=np.floataa  )  )
							self.assertEqual(examples['filename'][0]  ,      'my_name-train_29'  )
			def 					A_    (  self  ):
							_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(lowercase  ,      partial(dset.get_nearest_examples  ,      'vecs2'  ,      np.ones(5  ,      dtype=np.floataa  )  )  )






			def 					A_    (  self  ):
							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     :			Tuple    =    {'acknowledged': True}
											mocked_bulk.return_value([(True, None)] * 30  )
											_lowerCamelCase     :			List[Any]    =    {'hits': {'hits': [{'_score': 1, '_id': 29}]}}
											_lowerCamelCase     :			Optional[int]    =    Elasticsearch()

											dset.add_elasticsearch_index('filename'  ,      es_client=lowercase  )
											_lowerCamelCase, _lowerCamelCase     :			Union[str, Any]    =    dset.get_nearest_examples('filename'  ,      'my_name-train_29'  )
											self.assertEqual(examples['filename'][0]  ,      'my_name-train_29'  )






@require_faiss
class lowerCAmelCase__	(     lowercase						):






			'''simple docstring'''



			def 					A_    (  self  ):
							import faiss

							_lowerCamelCase     :			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
							_lowerCamelCase     :			Dict    =    np.zeros(5  ,      dtype=np.floataa  )
							_lowerCamelCase     :			Dict    =    1
							_lowerCamelCase, _lowerCamelCase     :			int    =    index.search(lowercase  )
							self.assertRaises(lowercase  ,      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(lowercase  )
							self.assertRaises(lowercase  ,      index.search_batch  ,      queries[0]  )
							_lowerCamelCase     :			List[str]    =    [scores[0] for scores in total_scores]
							_lowerCamelCase     :			List[Any]    =    [indices[0] for indices in total_indices]
							self.assertGreater(np.min(lowercase  )  ,      0  )
							self.assertListEqual([4, 3, 2, 1, 0]  ,      lowercase  )
			def 					A_    (  self  ):
							import faiss

							_lowerCamelCase     :			Tuple    =    FaissIndex(string_factory='Flat'  )
							index.add_vectors(np.eye(5  ,      dtype=np.floataa  )  )
							self.assertIsInstance(index.faiss_index  ,      faiss.IndexFlat  )
							_lowerCamelCase     :			int    =    FaissIndex(string_factory='LSH'  )
							index.add_vectors(np.eye(5  ,      dtype=np.floataa  )  )
							self.assertIsInstance(index.faiss_index  ,      faiss.IndexLSH  )
							with self.assertRaises(lowercase  ):
											_lowerCamelCase     :			List[Any]    =    FaissIndex(string_factory='Flat'  ,      custom_index=faiss.IndexFlat(5  )  )
			def 					A_    (  self  ):
							import faiss

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






			def 					A_    (  self  ):
							import faiss

							_lowerCamelCase     :			Tuple    =    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  )
											_lowerCamelCase     :			Union[str, Any]    =    FaissIndex.load(tmp_file.name  )
							os.unlink(tmp_file.name  )

							_lowerCamelCase     :			Tuple    =    np.zeros(5  ,      dtype=np.floataa  )
							_lowerCamelCase     :			Optional[int]    =    1
							_lowerCamelCase, _lowerCamelCase     :			Tuple    =    index.search(lowercase  )
							self.assertGreater(scores[0]  ,      0  )
							self.assertEqual(indices[0]  ,      1  )







@require_faiss
def      _snake_case			(				lowercase__		):
				import faiss

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

				_lowerCamelCase     :			Dict    =    'index.faiss'
				_lowerCamelCase     :			Optional[int]    =    f'''mock://{index_name}'''
				index.save(lowercase__							,    storage_options=mockfs.storage_options		)
				_lowerCamelCase     :			Dict    =    FaissIndex.load(lowercase__							,    storage_options=mockfs.storage_options		)

				_lowerCamelCase     :			Union[str, Any]    =    np.zeros(5							,    dtype=np.floataa		)
				_lowerCamelCase     :			Any    =    1
				_lowerCamelCase, _lowerCamelCase     :			List[Any]    =    index.search(lowercase__		)
				assert scores[0] > 0
				assert indices[0] == 1






@require_elasticsearch
class lowerCAmelCase__	(     lowercase						):






			'''simple docstring'''



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

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

											# single query with timeout
											_lowerCamelCase     :			List[str]    =    'foo'
											_lowerCamelCase     :			Union[str, Any]    =    {'hits': {'hits': [{'_score': 1, '_id': 0}]}}
											_lowerCamelCase, _lowerCamelCase     :			str    =    index.search(lowercase  ,      request_timeout=30  )
											self.assertEqual(scores[0]  ,      1  )
											self.assertEqual(indices[0]  ,      0  )

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

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





"""simple docstring"""






import json
import logging
import os
import sys
from time import time
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, require_torch_tpu


logging.basicConfig(level=logging.DEBUG)

lowercase__     						=						logging.getLogger()







def      _snake_case			(				lowercase__		):
				_lowerCamelCase     :			List[Any]    =    {}
				_lowerCamelCase     :			List[Any]    =    os.path.join(lowercase__							,    'all_results.json'		)
				if os.path.exists(lowercase__		):
								with open(lowercase__							,    'r'		) as f:
												_lowerCamelCase     :			List[Any]    =    json.load(lowercase__		)
				else:
								raise ValueError(f'''can\'t find {path}'''		)
				return results


lowercase__     						=						logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)






@require_torch_tpu
class lowerCAmelCase__	(     lowercase						):






			'''simple docstring'''



			def 					A_    (  self  ):
							import xla_spawn

							_lowerCamelCase     :			List[Any]    =    self.get_auto_remove_tmp_dir()
							_lowerCamelCase     :			List[Any]    =    F'''
            ./examples/pytorch/text-classification/run_glue.py
            --num_cores=8
            ./examples/pytorch/text-classification/run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --do_train
            --do_eval
            --debug tpu_metrics_debug
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --max_steps=10
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            '''.split()

							with patch.object(lowercase  ,      'argv'  ,      lowercase  ):
											_lowerCamelCase     :			Dict    =    time()
											xla_spawn.main()
											_lowerCamelCase     :			Any    =    time()

											_lowerCamelCase     :			Optional[int]    =    get_results(lowercase  )
											self.assertGreaterEqual(result['eval_accuracy']  ,      0.75  )

											# Assert that the script takes less than 500 seconds to make sure it doesn't hang.
											self.assertLess(end - start  ,      500  )






			def 					A_    (  self  ):
							import xla_spawn

							_lowerCamelCase     :			Tuple    =    '\n            ./tests/test_trainer_tpu.py\n            --num_cores=8\n            ./tests/test_trainer_tpu.py\n            '.split()
							with patch.object(lowercase  ,      'argv'  ,      lowercase  ):
											xla_spawn.main()