


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        _lowercase      (			unittest.TestCase       ):






				'''simple docstring'''



				def __init__(       self			:Dict			,			lowerCAmelCase__			:Dict			,			lowerCAmelCase__			:List[Any]=2			,			lowerCAmelCase__			:Optional[int]=56			,			lowerCAmelCase__			:Optional[Any]=True			,			lowerCAmelCase__			:Optional[Any]=True			,			lowerCAmelCase__			:List[Any]=True			,			lowerCAmelCase__			:Dict=True			,			lowerCAmelCase__			:int=99			,			lowerCAmelCase__			:str=32			,			lowerCAmelCase__			:int=2			,			lowerCAmelCase__			:Tuple=2			,			lowerCAmelCase__			:Any=7			,			lowerCAmelCase__			:Tuple="gelu_new"			,			lowerCAmelCase__			:str=0.1			,			lowerCAmelCase__			:List[Any]=0.1			,			lowerCAmelCase__			:Tuple=512			,			lowerCAmelCase__			:Optional[Any]=16			,			lowerCAmelCase__			:Optional[Any]=2			,			lowerCAmelCase__			:Optional[int]=0.02			,			lowerCAmelCase__			:int=4			,			lowerCAmelCase__			:Any="block_sparse"			,			lowerCAmelCase__			:Any=True			,			lowerCAmelCase__			:str=False			,			lowerCAmelCase__			:Optional[int]=2			,			lowerCAmelCase__			:Tuple=3			,			)		->			Tuple:
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		parent
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		batch_size
										__SCREAMING_SNAKE_CASE						:     Any					=		seq_length
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		is_training
										__SCREAMING_SNAKE_CASE						:     Tuple					=		use_attention_mask
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		use_token_type_ids
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		use_labels
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		vocab_size
										__SCREAMING_SNAKE_CASE						:     Tuple					=		hidden_size
										__SCREAMING_SNAKE_CASE						:     str					=		num_hidden_layers
										__SCREAMING_SNAKE_CASE						:     Tuple					=		num_attention_heads
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		intermediate_size
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		hidden_act
										__SCREAMING_SNAKE_CASE						:     str					=		hidden_dropout_prob
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		attention_probs_dropout_prob
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		max_position_embeddings
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		type_vocab_size
										__SCREAMING_SNAKE_CASE						:     Dict					=		type_sequence_label_size
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		initializer_range
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		num_choices

										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		rescale_embeddings
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		attention_type
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		use_bias
										__SCREAMING_SNAKE_CASE						:     Any					=		block_size
										__SCREAMING_SNAKE_CASE						:     Dict					=		num_random_blocks


				def 		__magic_name__(       self			:Tuple	)		->			List[str]:
										__SCREAMING_SNAKE_CASE						:     int					=		ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size	)

										__SCREAMING_SNAKE_CASE						:     Any					=		None
										if self.use_attention_mask:
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		random_attention_mask([self.batch_size, self.seq_length]	)

										__SCREAMING_SNAKE_CASE						:     Tuple					=		None
										if self.use_token_type_ids:
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		ids_tensor([self.batch_size, self.seq_length]			,			self.type_vocab_size	)

										__SCREAMING_SNAKE_CASE						:     Any					=		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=lowerCAmelCase__			,			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 		__magic_name__(       self			:List[Any]	)		->			Any:
										__SCREAMING_SNAKE_CASE						:     int					=		self.prepare_config_and_inputs()
										__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     int					=		config_and_inputs
										__SCREAMING_SNAKE_CASE						:     str					=		{
										    '''input_ids''': input_ids,
										    '''token_type_ids''': token_type_ids,
										    '''attention_mask''': attention_mask,
										}
										return config, inputs_dict




@require_flax
class        _lowercase      (			A__						,			unittest.TestCase       ):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			int						=		(
				    (
				        FlaxBigBirdForCausalLM,
				        FlaxBigBirdModel,
				        FlaxBigBirdForPreTraining,
				        FlaxBigBirdForMaskedLM,
				        FlaxBigBirdForMultipleChoice,
				        FlaxBigBirdForQuestionAnswering,
				        FlaxBigBirdForSequenceClassification,
				        FlaxBigBirdForTokenClassification,
				    )
				    if is_flax_available()
				    else ()
				)

				SCREAMING_SNAKE_CASE__				:			Optional[int]						=		False
				SCREAMING_SNAKE_CASE__				:			Optional[Any]						=		False


				def 		__magic_name__(       self			:Union[str, Any]	)		->			Dict:
										__SCREAMING_SNAKE_CASE						:     str					=		FlaxBigBirdModelTester(self	)


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


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


				@slow
				# copied from `test_modeling_flax_common` because it takes much longer than other models
				def 		__magic_name__(       self			:Dict	)		->			str:
										super().test_no_automatic_init()


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


				@slow
				def 		__magic_name__(       self			:Optional[Any]	)		->			Tuple:
										for model_class_name in self.all_model_classes:
																__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		model_class_name.from_pretrained('''google/bigbird-roberta-base'''	)
																self.assertIsNotNone(lowerCAmelCase__	)


				def 		__magic_name__(       self			:Union[str, Any]	)		->			List[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 		__magic_name__(       self			:Optional[Any]	)		->			str:
										__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     List[Any]					=		self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__	):
																						__SCREAMING_SNAKE_CASE						:     Any					=		self._prepare_for_class(lowerCAmelCase__			,			lowerCAmelCase__	)
																						__SCREAMING_SNAKE_CASE						:     str					=		model_class(lowerCAmelCase__	)

																						@jax.jit
																						def model_jitted(lowerCAmelCase__			:Union[str, Any]			,			lowerCAmelCase__			:Optional[int]=None			,			**lowerCAmelCase__			:Union[str, Any]	):
																												return model(input_ids=lowerCAmelCase__			,			attention_mask=lowerCAmelCase__			,			**lowerCAmelCase__	)

																						with self.subTest('''JIT Enabled'''	):
																												__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		model_jitted(**lowerCAmelCase__	).to_tuple()

																						with self.subTest('''JIT Disabled'''	):
																												with jax.disable_jit():
																																		__SCREAMING_SNAKE_CASE						:     Any					=		model_jitted(**lowerCAmelCase__	).to_tuple()

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



				def 		__magic_name__(       self			:Any			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:int			,			lowerCAmelCase__			:Dict			,			lowerCAmelCase__			:List[str]=1E-5			,			lowerCAmelCase__			:Optional[int]="outputs"			,			lowerCAmelCase__			:Optional[Any]=None	)		->			Tuple:
										# `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(lowerCAmelCase__			,			lowerCAmelCase__			,			lowerCAmelCase__			,			lowerCAmelCase__			,			lowerCAmelCase__			,			lowerCAmelCase__	)




import numpy as np
from scipy.spatial.distance import cdist
from sklearn.metrics import fa_score

import datasets


__lowerCAmelCase	: Optional[int]  				='\\n    @inproceedings{kakwani2020indicnlpsuite,\n    title={{IndicNLPSuite: Monolingual Corpora, Evaluation Benchmarks and Pre-trained Multilingual Language Models for Indian Languages}},\n    author={Divyanshu Kakwani and Anoop Kunchukuttan and Satish Golla and Gokul N.C. and Avik Bhattacharyya and Mitesh M. Khapra and Pratyush Kumar},\n    year={2020},\n    booktitle={Findings of EMNLP},\n}\n'

__lowerCAmelCase	: Optional[Any]  				='\\n    IndicGLUE is a natural language understanding benchmark for Indian languages. It contains a wide\n    variety of tasks and covers 11 major Indian languages - as, bn, gu, hi, kn, ml, mr, or, pa, ta, te.\n'

__lowerCAmelCase	: Dict  				='\nCompute IndicGLUE evaluation metric associated to each IndicGLUE dataset.\nArgs:\n    predictions: list of predictions to score (as int64),\n        except for \'cvit-mkb-clsr\' where each prediction is a vector (of float32).\n    references: list of ground truth labels corresponding to the predictions (as int64),\n        except for \'cvit-mkb-clsr\' where each reference is a vector (of float32).\nReturns: depending on the IndicGLUE subset, one or several of:\n    "accuracy": Accuracy\n    "f1": F1 score\n    "precision": Precision@10\nExamples:\n\n    >>> indic_glue_metric = datasets.load_metric(\'indic_glue\', \'wnli\')  # \'wnli\' or any of ["copa", "sna", "csqa", "wstp", "inltkh", "bbca", "iitp-mr", "iitp-pr", "actsa-sc", "md"]\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'accuracy\': 1.0}\n\n    >>> indic_glue_metric = datasets.load_metric(\'indic_glue\', \'wiki-ner\')\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'accuracy\': 1.0, \'f1\': 1.0}\n\n    >>> indic_glue_metric = datasets.load_metric(\'indic_glue\', \'cvit-mkb-clsr\')\n    >>> references = [[0.5, 0.5, 0.5], [0.1, 0.2, 0.3]]\n    >>> predictions = [[0.5, 0.5, 0.5], [0.1, 0.2, 0.3]]\n    >>> results = indic_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'precision@10\': 1.0}\n\n'


def        _UpperCamelCase							(      lowercase__    ,       lowercase__      ):
						return float((preds == labels).mean()      )


def        _UpperCamelCase							(      lowercase__    ,       lowercase__      ):
						__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		simple_accuracy(lowercase__    ,       lowercase__      )
						__SCREAMING_SNAKE_CASE						:     List[str]					=		float(fa_score(y_true=lowercase__    ,       y_pred=lowercase__      )      )
						return {
						    "accuracy": acc,
						    "f1": fa,
						}


def        _UpperCamelCase							(      lowercase__    ,       lowercase__      ):
						__SCREAMING_SNAKE_CASE						:     Optional[int]					=		np.array(lowercase__      )
						__SCREAMING_SNAKE_CASE						:     str					=		np.array(lowercase__      )
						__SCREAMING_SNAKE_CASE						:     str					=		en_sentvecs.shape[0]

						# mean centering
						__SCREAMING_SNAKE_CASE						:     Tuple					=		en_sentvecs - np.mean(lowercase__    ,       axis=0      )
						__SCREAMING_SNAKE_CASE						:     Optional[int]					=		in_sentvecs - np.mean(lowercase__    ,       axis=0      )

						__SCREAMING_SNAKE_CASE						:     str					=		cdist(lowercase__    ,       lowercase__    ,       '''cosine'''      )
						__SCREAMING_SNAKE_CASE						:     int					=		np.array(range(lowercase__      )      )
						__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		sim.argsort(axis=1      )[:, :10]
						__SCREAMING_SNAKE_CASE						:     str					=		np.any(preds == actual[:, None]    ,       axis=1      )
						return float(matches.mean()      )
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,			_KWARGS_DESCRIPTION       )
class        _lowercase      (			datasets.Metric       ):






				'''simple docstring'''



				def 		__magic_name__(       self			:Tuple	)		->			Tuple:
										if self.config_name not in [
										    "wnli",
										    "copa",
										    "sna",
										    "csqa",
										    "wstp",
										    "inltkh",
										    "bbca",
										    "cvit-mkb-clsr",
										    "iitp-mr",
										    "iitp-pr",
										    "actsa-sc",
										    "md",
										    "wiki-ner",
										]:
																raise KeyError(
																    '''You should supply a configuration name selected in '''
																    '''["wnli", "copa", "sna", "csqa", "wstp", "inltkh", "bbca", '''
																    '''"cvit-mkb-clsr", "iitp-mr", "iitp-pr", "actsa-sc", "md", '''
																    '''"wiki-ner"]'''	)
										return datasets.MetricInfo(
										    description=_DESCRIPTION			,			citation=_CITATION			,			inputs_description=_KWARGS_DESCRIPTION			,			features=datasets.Features(
										        {
										            '''predictions''': datasets.Value('''int64'''	)
										            if self.config_name != '''cvit-mkb-clsr'''
										            else datasets.Sequence(datasets.Value('''float32'''	)	),
										            '''references''': datasets.Value('''int64'''	)
										            if self.config_name != '''cvit-mkb-clsr'''
										            else datasets.Sequence(datasets.Value('''float32'''	)	),
										        }	)			,			codebase_urls=[]			,			reference_urls=[]			,			format='''numpy''' if self.config_name != '''cvit-mkb-clsr''' else None			,			)



				def 		__magic_name__(       self			:List[str]			,			lowerCAmelCase__			:Optional[Any]			,			lowerCAmelCase__			:Tuple	)		->			str:
										if self.config_name == "cvit-mkb-clsr":
																return {"precision@10": precision_at_aa(lowerCAmelCase__			,			lowerCAmelCase__	)}
										elif self.config_name in ["wiki-ner"]:
																return acc_and_fa(lowerCAmelCase__			,			lowerCAmelCase__	)
										elif self.config_name in [
										    "wnli",
										    "copa",
										    "sna",
										    "csqa",
										    "wstp",
										    "inltkh",
										    "bbca",
										    "iitp-mr",
										    "iitp-pr",
										    "actsa-sc",
										    "md",
										]:
																return {"accuracy": simple_accuracy(lowerCAmelCase__			,			lowerCAmelCase__	)}
										else:
																raise KeyError(
																    '''You should supply a configuration name selected in '''
																    '''["wnli", "copa", "sna", "csqa", "wstp", "inltkh", "bbca", '''
																    '''"cvit-mkb-clsr", "iitp-mr", "iitp-pr", "actsa-sc", "md", '''
																    '''"wiki-ner"]'''	)
