
import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
			import jax

			from transformers.models.big_bird.modeling_flax_big_bird import (
			    FlaxBigBirdForCausalLM,
			    FlaxBigBirdForMaskedLM,
			    FlaxBigBirdForMultipleChoice,
			    FlaxBigBirdForPreTraining,
			    FlaxBigBirdForQuestionAnswering,
			    FlaxBigBirdForSequenceClassification,
			    FlaxBigBirdForTokenClassification,
			    FlaxBigBirdModel,
			)




class     _SCREAMING_SNAKE_CASE     (      unittest.TestCase						):



	def __init__(    self						,  lowercase						,  lowercase=2						,  lowercase=56						,  lowercase=True						,  lowercase=True						,  lowercase=True						,  lowercase=True						,  lowercase=99						,  lowercase=32						,  lowercase=2						,  lowercase=2						,  lowercase=7						,  lowercase="gelu_new"						,  lowercase=0.1						,  lowercase=0.1						,  lowercase=512						,  lowercase=16						,  lowercase=2						,  lowercase=0.0_2						,  lowercase=4						,  lowercase="block_sparse"						,  lowercase=True						,  lowercase=False						,  lowercase=2						,  lowercase=3						,  )				->       List[Any]:
					lowerCamelCase_						=			parent
					lowerCamelCase_						=			batch_size
					lowerCamelCase_						=			seq_length
					lowerCamelCase_						=			is_training
					lowerCamelCase_						=			use_attention_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_choices

					lowerCamelCase_						=			rescale_embeddings
					lowerCamelCase_						=			attention_type
					lowerCamelCase_						=			use_bias
					lowerCamelCase_						=			block_size
					lowerCamelCase_						=			num_random_blocks



	def     SCREAMING_SNAKE_CASE_(    self )				->       int:
					lowerCamelCase_						=			ids_tensor([self.batch_size, self.seq_length]						,  self.vocab_size )

					lowerCamelCase_						=			None
					if self.use_attention_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_						=			BigBirdConfig(
					    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=lowercase						,  initializer_range=self.initializer_range						,  attention_type=self.attention_type						,  block_size=self.block_size						,  num_random_blocks=self.num_random_blocks						,  use_bias=self.use_bias						,  rescale_embeddings=self.rescale_embeddings						,  )

					return config, input_ids, token_type_ids, attention_mask



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					lowerCamelCase_						=			self.prepare_config_and_inputs()
					lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_						=			config_and_inputs
					lowerCamelCase_						=			{
					    "input_ids": input_ids,
					    "token_type_ids": token_type_ids,
					    "attention_mask": attention_mask,
					}
					return config, inputs_dict




@require_flax
class     _SCREAMING_SNAKE_CASE     (      snake_case_			,       unittest.TestCase						):
	lowerCAmelCase__         = (
	    (
	        FlaxBigBirdForCausalLM,
	        FlaxBigBirdModel,
	        FlaxBigBirdForPreTraining,
	        FlaxBigBirdForMaskedLM,
	        FlaxBigBirdForMultipleChoice,
	        FlaxBigBirdForQuestionAnswering,
	        FlaxBigBirdForSequenceClassification,
	        FlaxBigBirdForTokenClassification,
	    )
	    if is_flax_available()
	    else ()
	)

	lowerCAmelCase__         = False
	lowerCAmelCase__         = False



	def     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					lowerCamelCase_						=			FlaxBigBirdModelTester(self )



	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[int]:
					super().test_from_pretrained_save_pretrained()



	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					super().test_from_pretrained_with_no_automatic_init()



	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					super().test_no_automatic_init()



	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def     SCREAMING_SNAKE_CASE_(    self )				->       int:
					super().test_hidden_states_output()



	@slow
	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					for model_class_name in self.all_model_classes:
									lowerCamelCase_						=			model_class_name.from_pretrained("google/bigbird-roberta-base" )
									self.assertIsNotNone(lowercase )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					if self.test_attn_probs:
									super().test_attention_outputs()



	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[int]:
					lowerCamelCase_	,    lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									with self.subTest(model_class.__name__ ):
													lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase )
													lowerCamelCase_						=			model_class(lowercase )

													@jax.jit
													def model_jitted(lowercase						,  lowercase=None						,  **lowercase ):
																	return model(input_ids=lowercase						,  attention_mask=lowercase						,  **lowercase )

													with self.subTest("JIT Enabled" ):
																	lowerCamelCase_						=			model_jitted(**lowercase ).to_tuple()

													with self.subTest("JIT Disabled" ):
																	with jax.disable_jit():
																					lowerCamelCase_						=			model_jitted(**lowercase ).to_tuple()

													self.assertEqual(len(lowercase )						,  len(lowercase ) )
													for jitted_output, output in zip(lowercase						,  lowercase ):
																	self.assertEqual(jitted_output.shape						,  output.shape )



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase						,  lowercase						,  lowercase=1e-5						,  lowercase="outputs"						,  lowercase=None )				->       Optional[int]:
					# `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
					# an effort was done to return `attention_probs` (yet to be verified).
					if name.startswith("outputs.attentions" ):
									return
					else:
									super().check_pt_flax_outputs(lowercase						,  lowercase						,  lowercase						,  lowercase						,  lowercase						,  lowercase )




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				=pytest.mark.integration




@require_faiss
class     _SCREAMING_SNAKE_CASE     (      snake_case_						):



	def     SCREAMING_SNAKE_CASE_(    self )				->       List[str]:
					lowerCamelCase_						=			Dataset.from_dict({"filename": ["my_name-train" + "_" + str(lowercase ) for x in np.arange(30 ).tolist()]} )
					return dset



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					import faiss

					lowerCamelCase_						=			self._create_dummy_dataset()
					lowerCamelCase_						=			dset.map(
					    lambda lowercase						,  lowercase : {"vecs": i * np.ones(5						,  dtype=np.floataa )}						,  with_indices=lowercase						,  keep_in_memory=lowercase )
					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     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					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     SCREAMING_SNAKE_CASE_(    self )				->       Optional[int]:
					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=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_						=			dset.get_nearest_examples("vecs2"						,  np.ones(5						,  dtype=np.floataa ) )
					self.assertEqual(examples["filename"][0]						,  "my_name-train_29" )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					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(lowercase						,  partial(dset.get_nearest_examples						,  "vecs2"						,  np.ones(5						,  dtype=np.floataa ) ) )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[int]:
					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=lowercase )
									lowerCamelCase_	,    lowerCamelCase_						=			dset.get_nearest_examples("filename"						,  "my_name-train_29" )
									self.assertEqual(examples["filename"][0]						,  "my_name-train_29" )




@require_faiss
class     _SCREAMING_SNAKE_CASE     (      snake_case_						):



	def     SCREAMING_SNAKE_CASE_(    self )				->       Tuple:
					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(lowercase )
					self.assertRaises(lowercase						,  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(lowercase )
					self.assertRaises(lowercase						,  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(lowercase )						,  0 )
					self.assertListEqual([4, 3, 2, 1, 0]						,  lowercase )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Any:
					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(lowercase ):
									lowerCamelCase_						=			FaissIndex(string_factory="Flat"						,  custom_index=faiss.IndexFlat(5 ) )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[int]:
					import faiss

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



	def     SCREAMING_SNAKE_CASE_(    self )				->       List[str]:
					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=lowercase ) 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(lowercase )
					self.assertGreater(scores[0]						,  0 )
					self.assertEqual(indices[0]						,  1 )





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

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





@require_elasticsearch
class     _SCREAMING_SNAKE_CASE     (      snake_case_						):



	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[Any]:
					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=lowercase )
									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(lowercase )
									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(lowercase						,  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(lowercase )
									lowerCamelCase_						=			[scores[0] for scores in total_scores]
									lowerCamelCase_						=			[indices[0] for indices in total_indices]
									self.assertGreater(np.min(lowercase )						,  0 )
									self.assertListEqual([1, 1, 1]						,  lowercase )

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


