



"""simple docstring"""




import math
import unittest

from transformers import BioGptConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    BioGptForCausalLM,
					    BioGptForSequenceClassification,
					    BioGptForTokenClassification,
					    BioGptModel,
					    BioGptTokenizer,
					)
					from transformers.models.biogpt.modeling_biogpt import BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST






class 					__magic_name__							:






				'''simple docstring'''





				def __init__( self	,				_a	,				_a=13	,				_a=7	,				_a=True	,				_a=True	,				_a=False	,				_a=True	,				_a=99	,				_a=32	,				_a=5	,				_a=4	,				_a=37	,				_a="gelu"	,				_a=0.1	,				_a=0.1	,				_a=512	,				_a=16	,				_a=2	,				_a=0.02	,				_a=3	,				_a=4	,				_a=None	,				):





							"""simple docstring"""
							lowerCamelCase       =      parent
							lowerCamelCase       =      batch_size
							lowerCamelCase       =      seq_length
							lowerCamelCase       =      is_training
							lowerCamelCase       =      use_input_mask
							lowerCamelCase       =      use_token_type_ids
							lowerCamelCase       =      use_labels
							lowerCamelCase       =      vocab_size
							lowerCamelCase       =      hidden_size
							lowerCamelCase       =      num_hidden_layers
							lowerCamelCase       =      num_attention_heads
							lowerCamelCase       =      intermediate_size
							lowerCamelCase       =      hidden_act
							lowerCamelCase       =      hidden_dropout_prob
							lowerCamelCase       =      attention_probs_dropout_prob
							lowerCamelCase       =      max_position_embeddings
							lowerCamelCase       =      type_vocab_size
							lowerCamelCase       =      type_sequence_label_size
							lowerCamelCase       =      initializer_range
							lowerCamelCase       =      num_labels
							lowerCamelCase       =      num_choices
							lowerCamelCase       =      scope





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      ids_tensor([self.batch_size, self.seq_length]	,				self.vocab_size				)

							lowerCamelCase       =      None
							if self.use_input_mask:
										lowerCamelCase       =      random_attention_mask([self.batch_size, self.seq_length]				)

							lowerCamelCase       =      None
							if self.use_token_type_ids:
										lowerCamelCase       =      ids_tensor([self.batch_size, self.seq_length]	,				self.type_vocab_size				)

							lowerCamelCase       =      None
							lowerCamelCase       =      None
							lowerCamelCase       =      None
							if self.use_labels:
										lowerCamelCase       =      ids_tensor([self.batch_size]	,				self.type_sequence_label_size				)
										lowerCamelCase       =      ids_tensor([self.batch_size, self.seq_length]	,				self.num_labels				)
										lowerCamelCase       =      ids_tensor([self.batch_size]	,				self.num_choices				)

							lowerCamelCase       =      self.get_config()

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return BioGptConfig(
							    vocab_size=self.vocab_size	,				hidden_size=self.hidden_size	,				num_hidden_layers=self.num_hidden_layers	,				num_attention_heads=self.num_attention_heads	,				intermediate_size=self.intermediate_size	,				hidden_act=self.hidden_act	,				hidden_dropout_prob=self.hidden_dropout_prob	,				attention_probs_dropout_prob=self.attention_probs_dropout_prob	,				max_position_embeddings=self.max_position_embeddings	,				type_vocab_size=self.type_vocab_size	,				is_decoder=_A	,				initializer_range=self.initializer_range	,				)





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a	,				_a	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptModel(config=_A				)
							model.to(_A				)
							model.eval()
							lowerCamelCase       =      model(_A	,				attention_mask=_A				)
							lowerCamelCase       =      model(_A				)
							self.parent.assertEqual(result.last_hidden_state.shape	,				(self.batch_size, self.seq_length, self.hidden_size)				)





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a	,				_a	,				_a	,				_a	,				_a	,				_a	,				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptForCausalLM(config=_A				)
							model.to(_A				)
							model.eval()
							lowerCamelCase       =      model(_A	,				attention_mask=_A	,				token_type_ids=_A	,				labels=_A				)
							self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.seq_length, self.vocab_size)				)





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a	,				_a	,				*_a				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptModel(config=_A				)
							model.to(_A				)
							model.eval()

							# create attention mask
							lowerCamelCase       =      torch.ones(input_ids.shape	,				dtype=torch.long	,				device=_A				)
							lowerCamelCase       =      self.seq_length // 2
							lowerCamelCase       =      0

							# first forward pass
							lowerCamelCase       =      model(_A	,				attention_mask=_A				).to_tuple()

							# create hypothetical next token and extent to next_input_ids
							lowerCamelCase       =      ids_tensor((self.batch_size, 1)	,				config.vocab_size				)

							# change a random masked slice from input_ids
							lowerCamelCase       =      ids_tensor((1,)	,				_A				).item() + 1
							lowerCamelCase       =      ids_tensor((self.batch_size, 1)	,				config.vocab_size				).squeeze(-1				)
							lowerCamelCase       =      random_other_next_tokens

							# append to next input_ids and attn_mask
							lowerCamelCase       =      torch.cat([input_ids, next_tokens]	,				dim=-1				)
							lowerCamelCase       =      torch.cat(
							    [attn_mask, torch.ones((attn_mask.shape[0], 1)	,				dtype=torch.long	,				device=_A				)]	,				dim=1	,				)

							# get two different outputs
							lowerCamelCase       =      model(_A	,				attention_mask=_A				)['last_hidden_state']
							lowerCamelCase       =      model(_A	,				past_key_values=_A	,				attention_mask=_A				)['last_hidden_state']

							# select random slice
							lowerCamelCase       =      ids_tensor((1,)	,				output_from_past.shape[-1]				).item()
							lowerCamelCase       =      output_from_no_past[:, -1, random_slice_idx].detach()
							lowerCamelCase       =      output_from_past[:, 0, random_slice_idx].detach()

							# test that outputs are equal for slice
							self.parent.assertTrue(torch.allclose(_A	,				_A	,				atol=1e-3				)				)





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a	,				_a	,				*_a				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptModel(config=_A				).to(_A				).eval()

							lowerCamelCase       =      torch.ones(input_ids.shape	,				dtype=torch.long	,				device=_A				)

							# first forward pass
							lowerCamelCase       =      model(_A	,				attention_mask=_A	,				use_cache=_A				)

							lowerCamelCase       =      outputs.to_tuple()

							# create hypothetical multiple next token and extent to next_input_ids
							lowerCamelCase       =      ids_tensor((self.batch_size, 3)	,				config.vocab_size				)
							lowerCamelCase       =      ids_tensor((self.batch_size, 3)	,				2				)

							# append to next input_ids and
							lowerCamelCase       =      torch.cat([input_ids, next_tokens]	,				dim=-1				)
							lowerCamelCase       =      torch.cat([attention_mask, next_attn_mask]	,				dim=-1				)

							lowerCamelCase       =      model(_A	,				attention_mask=_A				)['last_hidden_state']
							lowerCamelCase       =      model(_A	,				attention_mask=_A	,				past_key_values=_A				)[
							    'last_hidden_state'
							]

							# select random slice
							lowerCamelCase       =      ids_tensor((1,)	,				output_from_past.shape[-1]				).item()
							lowerCamelCase       =      output_from_no_past[:, -3:, random_slice_idx].detach()
							lowerCamelCase       =      output_from_past[:, :, random_slice_idx].detach()

							self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]				)

							# test that outputs are equal for slice
							self.parent.assertTrue(torch.allclose(_A	,				_A	,				atol=1e-3				)				)





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a	,				_a	,				*_a	,				_a=False				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptForCausalLM(_A				)
							model.to(_A				)
							if gradient_checkpointing:
										model.gradient_checkpointing_enable()

							lowerCamelCase       =      model(_A	,				labels=_A				)
							self.parent.assertEqual(result.loss.shape	,				()				)
							self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.seq_length, self.vocab_size)				)
							result.loss.backward()





				def      _lowerCAmelCase       ( self	,				_a	,				*_a				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptModel(_A				)
							lowerCamelCase       =      model.config.initializer_range / math.sqrt(2 * model.config.num_hidden_layers				)
							for key in model.state_dict().keys():
										if "c_proj" in key and "weight" in key:
													self.parent.assertLessEqual(abs(torch.std(model.state_dict()[key]				) - model_std				)	,				0.001				)
													self.parent.assertLessEqual(abs(torch.mean(model.state_dict()[key]				) - 0.0				)	,				0.01				)





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a	,				_a	,				*_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.num_labels
							lowerCamelCase       =      BioGptForTokenClassification(_A				)
							model.to(_A				)
							model.eval()
							lowerCamelCase       =      model(_A	,				attention_mask=_A	,				token_type_ids=_A				)
							self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.seq_length, self.num_labels)				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.prepare_config_and_inputs()
							(
							    lowerCamelCase       
							)       =      config_and_inputs
							lowerCamelCase       =      {'input_ids': input_ids, 'attention_mask': input_mask}
							return config, inputs_dict






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






				'''simple docstring'''
				__UpperCamelCase    						=		(
				    (BioGptModel, BioGptForCausalLM, BioGptForSequenceClassification, BioGptForTokenClassification)
				    if is_torch_available()
				    else ()
				)
				__UpperCamelCase    						=		(BioGptForCausalLM,) if is_torch_available() else ()
				__UpperCamelCase    						=		(
				    {
				        '''feature-extraction''': BioGptModel,
				        '''text-classification''': BioGptForSequenceClassification,
				        '''text-generation''': BioGptForCausalLM,
				        '''token-classification''': BioGptForTokenClassification,
				        '''zero-shot''': BioGptForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__UpperCamelCase    						=		False





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptModelTester(self				)
							lowerCamelCase       =      ConfigTester(self	,				config_class=_A	,				hidden_size=37				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							self.config_tester.run_common_tests()





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_A				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
										lowerCamelCase       =      type
										self.model_tester.create_and_check_model(*_A				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_biogpt_model_attention_mask_past(*_A				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_forward_and_backwards(*_A	,				gradient_checkpointing=_A				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_biogpt_model_past_large_inputs(*_A				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_biogpt_weight_initialization(*_A				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_biogpt_for_token_classification(*_A				)





				@slow
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptForCausalLM.from_pretrained("""microsoft/biogpt"""				)
							model.to(_A				)
							lowerCamelCase       =      BioGptTokenizer.from_pretrained("""microsoft/biogpt"""				)

							lowerCamelCase       =      'left'

							# Define PAD Token = EOS Token = 50256
							lowerCamelCase       =      tokenizer.eos_token
							lowerCamelCase       =      model.config.eos_token_id

							# use different length sentences to test batching
							lowerCamelCase       =      [
							    'Hello, my dog is a little',
							    'Today, I',
							]

							lowerCamelCase       =      tokenizer(_A	,				return_tensors="""pt"""	,				padding=_A				)
							lowerCamelCase       =      inputs['input_ids'].to(_A				)

							lowerCamelCase       =      model.generate(
							    input_ids=_A	,				attention_mask=inputs["""attention_mask"""].to(_A				)	,				)

							lowerCamelCase       =      tokenizer(sentences[0]	,				return_tensors="""pt"""				).input_ids.to(_A				)
							lowerCamelCase       =      model.generate(input_ids=_A				)

							lowerCamelCase       =      inputs_non_padded.shape[-1] - inputs['attention_mask'][-1].long().sum().cpu().item()
							lowerCamelCase       =      tokenizer(sentences[1]	,				return_tensors="""pt"""				).input_ids.to(_A				)
							lowerCamelCase       =      model.generate(input_ids=_A	,				max_length=model.config.max_length - num_paddings				)

							lowerCamelCase       =      tokenizer.batch_decode(_A	,				skip_special_tokens=_A				)
							lowerCamelCase       =      tokenizer.decode(output_non_padded[0]	,				skip_special_tokens=_A				)
							lowerCamelCase       =      tokenizer.decode(output_padded[0]	,				skip_special_tokens=_A				)

							lowerCamelCase       =      [
							    'Hello, my dog is a little bit bigger than a little bit.',
							    'Today, I have a good idea of how to use the information',
							]
							self.assertListEqual(_A	,				_A				)
							self.assertListEqual(_A	,				[non_padded_sentence, padded_sentence]				)





				@slow
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							for model_name in BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCamelCase       =      BioGptModel.from_pretrained(_A				)
										self.assertIsNotNone(_A				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs_for_common()
							lowerCamelCase       =      3
							lowerCamelCase       =      input_dict['input_ids']
							lowerCamelCase       =      input_ids.ne(1				).to(_A				)
							lowerCamelCase       =      ids_tensor([self.model_tester.batch_size]	,				self.model_tester.type_sequence_label_size				)
							lowerCamelCase       =      BioGptForSequenceClassification(_A				)
							model.to(_A				)
							model.eval()
							lowerCamelCase       =      model(_A	,				attention_mask=_A	,				labels=_A				)
							self.assertEqual(result.logits.shape	,				(self.model_tester.batch_size, self.model_tester.num_labels)				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs_for_common()
							lowerCamelCase       =      3
							lowerCamelCase       =      'multi_label_classification'
							lowerCamelCase       =      input_dict['input_ids']
							lowerCamelCase       =      input_ids.ne(1				).to(_A				)
							lowerCamelCase       =      ids_tensor(
							    [self.model_tester.batch_size, config.num_labels]	,				self.model_tester.type_sequence_label_size				).to(torch.float				)
							lowerCamelCase       =      BioGptForSequenceClassification(_A				)
							model.to(_A				)
							model.eval()
							lowerCamelCase       =      model(_A	,				attention_mask=_A	,				labels=_A				)
							self.assertEqual(result.logits.shape	,				(self.model_tester.batch_size, self.model_tester.num_labels)				)






@require_torch
class 					__magic_name__							( unittest.TestCase       ):






				'''simple docstring'''





				@slow
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptForCausalLM.from_pretrained("""microsoft/biogpt"""				)
							lowerCamelCase       =      torch.tensor([[2, 4_805, 9, 656, 21]]				)
							lowerCamelCase       =      model(_A				)[0]

							lowerCamelCase       =      42_384

							lowerCamelCase       =      torch.Size((1, 5, vocab_size)				)
							self.assertEqual(output.shape	,				_A				)

							lowerCamelCase       =      torch.tensor(
							    [[[-9.5_236, -9.8_918, 10.4_557], [-11.0_469, -9.6_423, 8.1_022], [-8.8_664, -7.8_826, 5.5_325]]]				)

							self.assertTrue(torch.allclose(output[:, :3, :3]	,				_A	,				atol=1e-4				)				)





				@slow
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      BioGptTokenizer.from_pretrained("""microsoft/biogpt"""				)
							lowerCamelCase       =      BioGptForCausalLM.from_pretrained("""microsoft/biogpt"""				)
							model.to(_A				)

							torch.manual_seed(0				)
							lowerCamelCase       =      tokenizer("""COVID-19 is"""	,				return_tensors="""pt"""				).to(_A				)
							lowerCamelCase       =      model.generate(
							    **_A	,				min_length=100	,				max_length=1_024	,				num_beams=5	,				early_stopping=_A	,				)
							lowerCamelCase       =      tokenizer.decode(output_ids[0]	,				skip_special_tokens=_A				)

							lowerCamelCase       =      (
							    'COVID-19 is a global pandemic caused by severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2), the'
							    ' causative agent of coronavirus disease 2019 (COVID-19), which has spread to more than 200 countries and'
							    ' territories, including the United States (US), Canada, Australia, New Zealand, the United Kingdom (UK),'
							    ' and the United States of America (USA), as of March 11, 2020, with more than 800,000 confirmed cases and'
							    ' more than 800,000 deaths.'
							)
							self.assertEqual(_A	,				_A				)









"""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							:    List[Any]  							=							pytest.mark.integration






@require_faiss
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''





				def      _lowerCAmelCase       ( self				):





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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

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





							"""simple docstring"""
							import faiss

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      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(_a	,				partial(dset.get_nearest_examples	,				"""vecs2"""	,				np.ones(5	,				dtype=np.floataa				)				)				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							from elasticsearch import Elasticsearch

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

										dset.add_elasticsearch_index("""filename"""	,				es_client=_a				)
										lowerCamelCase   ,     lowerCamelCase       =      dset.get_nearest_examples("""filename"""	,				"""my_name-train_29"""				)
										self.assertEqual(examples["""filename"""][0]	,				"""my_name-train_29"""				)






@require_faiss
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

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

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

							lowerCamelCase       =      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=_a				) as tmp_file:
										index.save(tmp_file.name				)
										lowerCamelCase       =      FaissIndex.load(tmp_file.name				)
							os.unlink(tmp_file.name				)

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




@require_faiss
def 					a__				(       snake_case__       )		->      Tuple:
			import faiss

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

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

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






@require_elasticsearch
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''





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

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

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

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

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



