"""simple docstring"""







import jax
import jax.numpy as jnp
from bigbird_flax import FlaxBigBirdForNaturalQuestions
from datasets import load_from_disk

from transformers import BigBirdTokenizerFast


SCREAMING_SNAKE_CASE						=    {0: 'null', 1: 'short', 2: 'long', 3: 'yes', 4: 'no'}
SCREAMING_SNAKE_CASE						=    set(''.join(['‘', '’', '´', '`', '.', ',', '-', '"']))


def   __lowerCAmelCase(		__UpperCAmelCase  ,__UpperCAmelCase=0  ,__UpperCAmelCase=None      ):






	"""simple docstring"""





	return [" ".join(x.split(' '      )[begin:end]      ) for x in answers if len(x.split(' '      )      ) > 1]


def   __lowerCAmelCase(		__UpperCAmelCase  ,__UpperCAmelCase=False      ):






	"""simple docstring"""





	if make_sub_answers:
		# if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
		# *e.g.* if the correct answer contains a prefix such as "the", or "a"
		_lowercase			:     Optional[Any]      					=				(
		    given_answers + get_sub_answers(__UpperCAmelCase  ,begin=1      ) + get_sub_answers(__UpperCAmelCase  ,end=-1      )
		)
	_lowercase			:     str      					=				[]
	for answer in given_answers:
		_lowercase			:     List[str]      					=				answer.replace('_'  ,' '      ).lower()
		_lowercase			:     int      					=				''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ' ' for c in alias      )
		answers.append(' '.join(alias.split()      ).strip()      )
	return set(__UpperCAmelCase      )


def   __lowerCAmelCase(		__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase=1  ,__UpperCAmelCase=100      ):






	"""simple docstring"""





	_lowercase					,   _lowercase			:     Optional[Any]      					=				jax.lax.top_k(__UpperCAmelCase  ,__UpperCAmelCase      )
	_lowercase					,   _lowercase			:     Optional[int]      					=				jax.lax.top_k(__UpperCAmelCase  ,__UpperCAmelCase      )

	_lowercase			:     str      					=				best_end_idx[:, None] - best_start_idx[None, :]
	_lowercase			:     List[str]      					=				jnp.logical_or(widths < 0  ,widths > max_size      )
	_lowercase			:     int      					=				(best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
	_lowercase			:     str      					=				jnp.argmax(__UpperCAmelCase      ).item()

	return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]


def   __lowerCAmelCase(		__UpperCAmelCase      ):






	"""simple docstring"""





	_lowercase			:     str      					=				sample['question']['text']
	_lowercase			:     Tuple      					=				sample['document']['tokens']['token']
	_lowercase			:     Union[str, Any]      					=				sample['document']['tokens']['is_html']
	_lowercase			:     List[str]      					=				sample['annotations']['long_answer']
	_lowercase			:     List[str]      					=				sample['annotations']['short_answers']

	_lowercase			:     Optional[int]      					=				' '.join([context[i] for i in range(len(__UpperCAmelCase      )      ) if not is_html[i]]      )

	# 0 - No ; 1 - Yes
	for answer in sample["annotations"]["yes_no_answer"]:
		if answer == 0 or answer == 1:
			return {
			    "question": question,
			    "context": context_string,
			    "short": [],
			    "long": [],
			    "category": "no" if answer == 0 else "yes",
			}

	_lowercase			:     str      					=				[]
	for s in short_answers:
		short_targets.extend(s['text']      )
	_lowercase			:     List[Any]      					=				list(set(__UpperCAmelCase      )      )

	_lowercase			:     Optional[int]      					=				[]
	for s in long_answers:
		if s["start_token"] == -1:
			continue
		_lowercase			:     str      					=				context[s['start_token'] : s['end_token']]
		_lowercase			:     str      					=				is_html[s['start_token'] : s['end_token']]
		_lowercase			:     List[Any]      					=				' '.join([answer[i] for i in range(len(__UpperCAmelCase      )      ) if not html[i]]      )
		if new_answer not in long_targets:
			long_targets.append(__UpperCAmelCase      )

	_lowercase			:     Optional[int]      					=				'long_short' if len(short_targets + long_targets      ) > 0 else 'null'

	return {
	    "question": question,
	    "context": context_string,
	    "short": short_targets,
	    "long": long_targets,
	    "category": category,
	}


def   __lowerCAmelCase(		):






	"""simple docstring"""





	_lowercase			:     List[str]      					=				load_from_disk('natural-questions-validation'      )
	_lowercase			:     int      					=				dataset.map(__UpperCAmelCase      ).remove_columns(['annotations', 'document', 'id']      )
	print(__UpperCAmelCase      )

	_lowercase			:     Optional[int]      					=				dataset.filter(lambda __UpperCAmelCase      : (len(x['question']      ) + len(x['context']      )) < 4 * 4_096      )
	_lowercase			:     str      					=				short_validation_dataset.filter(lambda __UpperCAmelCase      : x["category"] != "null"      )
	short_validation_dataset

	_lowercase			:     Union[str, Any]      					=				'vasudevgupta/flax-bigbird-natural-questions'
	_lowercase			:     Dict      					=				FlaxBigBirdForNaturalQuestions.from_pretrained(__UpperCAmelCase      )
	_lowercase			:     str      					=				BigBirdTokenizerFast.from_pretrained(__UpperCAmelCase      )

	@jax.jit
	def forward(*__UpperCAmelCase  ,**__UpperCAmelCase      ):
		_lowercase					,   _lowercase					,   _lowercase			:     List[str]      					=				model(*__UpperCAmelCase  ,**__UpperCAmelCase      )
		return start_logits, end_logits, jnp.argmax(__UpperCAmelCase  ,axis=-1      )

	def evaluate(__UpperCAmelCase      ):
		# encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
		_lowercase			:     Any      					=				tokenizer(
		    example['question']  ,example['context']  ,return_tensors='np'  ,max_length=4_096  ,padding='max_length'  ,truncation=__UpperCAmelCase  ,)

		_lowercase					,   _lowercase					,   _lowercase			:     List[str]      					=				forward(**__UpperCAmelCase      )

		_lowercase			:     Any      					=				CATEGORY_MAPPING[category.item()]

		_lowercase			:     str      					=				example['long'] + example['short']
		if example["category"] in ["yes", "no", "null"]:
			_lowercase			:     str      					=				[example['category']]
		_lowercase			:     Optional[Any]      					=				example['category'] != 'null'
		# Now target can be: "yes", "no", "null", "list of long & short answers"

		if predicted_category in ["yes", "no", "null"]:
			_lowercase			:     Any      					=				[predicted_category]
			_lowercase			:     List[str]      					=				example['output'] == example['targets']
			_lowercase			:     str      					=				predicted_category != 'null'
			return example

		_lowercase			:     Optional[int]      					=				38 if predicted_category == 'short' else 1_024
		_lowercase					,   _lowercase			:     Optional[Any]      					=				get_best_valid_start_end_idx(
		    start_scores[0]  ,end_scores[0]  ,top_k=8  ,max_size=__UpperCAmelCase      )

		_lowercase			:     List[str]      					=				inputs['input_ids'][0].tolist()
		_lowercase			:     List[Any]      					=				[tokenizer.decode(input_ids[start_score : end_score + 1]      )]

		_lowercase			:     List[str]      					=				expand_to_aliases(example['targets']  ,make_sub_answers=__UpperCAmelCase      )
		_lowercase			:     Union[str, Any]      					=				expand_to_aliases(example['output']      )

		# some preprocessing to both prediction and answer
		_lowercase			:     str      					=				{''.join(a.split()      ) for a in answers}
		_lowercase			:     int      					=				{''.join(p.split()      ) for p in predictions}
		_lowercase			:     int      					=				{s for s in predictions if s not in ['``', '\'\'', '`', '\'']}

		# if there is a common element, it's a exact match
		_lowercase			:     Dict      					=				len(list(answers & predictions      )      ) > 0
		_lowercase			:     Optional[int]      					=				predicted_category != 'null' and len(__UpperCAmelCase      ) > 0

		return example

	_lowercase			:     Any      					=				short_validation_dataset.map(__UpperCAmelCase      )

	_lowercase			:     int      					=				len(__UpperCAmelCase      )
	_lowercase			:     List[str]      					=				len(short_validation_dataset.filter(lambda __UpperCAmelCase      : x["match"] == 1      )      )
	print('EM score:'  ,(matched / total) * 100  ,'%'      )


if __name__ == "__main__":
						main()









from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

		from transformers import (
		    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFXLMForMultipleChoice,
		    TFXLMForQuestionAnsweringSimple,
		    TFXLMForSequenceClassification,
		    TFXLMForTokenClassification,
		    TFXLMModel,
		    TFXLMWithLMHeadModel,
		    XLMConfig,
		)




class        snake_case__				:
				"""simple docstring"""



				def __init__( self							:	Optional[Any]   , __lowerCamelCase							:	Any   , )   ->	Any:
							a       =     parent
							a       =     13
							a       =     7
							a       =     True
							a       =     True
							a       =     True
							a       =     True
							a       =     True
							a       =     False
							a       =     False
							a       =     False
							a       =     2
							a       =     99
							a       =     0
							a       =     32
							a       =     2
							a       =     4
							a       =     0.1
							a       =     0.1
							a       =     5_12
							a       =     16
							a       =     2
							a       =     0.02
							a       =     3
							a       =     4
							a       =     "last"
							a       =     True
							a       =     None
							a       =     0





				def 				__UpperCAmelCase				( self							:	int				)   ->	Tuple:
							a       =     ids_tensor([self.batch_size, self.seq_length]   , self.vocab_size				)
							a       =     random_attention_mask([self.batch_size, self.seq_length]   , dtype=tf.floataa				)

							a       =     None
							if self.use_input_lengths:
										a       =     (
										    ids_tensor([self.batch_size]   , vocab_size=2				) + self.seq_length - 2
										)  # small variation of seq_length

							a       =     None
							if self.use_token_type_ids:
										a       =     ids_tensor([self.batch_size, self.seq_length]   , self.n_langs				)

							a       =     None
							a       =     None
							a       =     None
							if self.use_labels:
										a       =     ids_tensor([self.batch_size]   , self.type_sequence_label_size				)
										a       =     ids_tensor([self.batch_size, self.seq_length]   , self.num_labels				)
										a       =     ids_tensor([self.batch_size]   , 2   , dtype=tf.floataa				)
										a       =     ids_tensor([self.batch_size]   , self.num_choices				)

							a       =     XLMConfig(
							    vocab_size=self.vocab_size   , n_special=self.n_special   , emb_dim=self.hidden_size   , n_layers=self.num_hidden_layers   , n_heads=self.num_attention_heads   , dropout=self.hidden_dropout_prob   , attention_dropout=self.attention_probs_dropout_prob   , gelu_activation=self.gelu_activation   , sinusoidal_embeddings=self.sinusoidal_embeddings   , asm=self.asm   , causal=self.causal   , n_langs=self.n_langs   , max_position_embeddings=self.max_position_embeddings   , initializer_range=self.initializer_range   , summary_type=self.summary_type   , use_proj=self.use_proj   , bos_token_id=self.bos_token_id   , )

							return (
							    config,
							    input_ids,
							    token_type_ids,
							    input_lengths,
							    sequence_labels,
							    token_labels,
							    is_impossible_labels,
							    choice_labels,
							    input_mask,
							)





				def 				__UpperCAmelCase				( self							:	List[str]   , __lowerCamelCase							:	List[str]   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	List[Any]   , __lowerCamelCase							:	List[Any]   , __lowerCamelCase							:	List[Any]   , __lowerCamelCase							:	List[Any]   , __lowerCamelCase							:	Any   , )   ->	Union[str, Any]:
							a       =     TFXLMModel(config=__lowerCamelCase				)
							a       =     {"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
							a       =     model(__lowerCamelCase				)

							a       =     [input_ids, input_mask]
							a       =     model(__lowerCamelCase				)
							self.parent.assertEqual(result.last_hidden_state.shape   , (self.batch_size, self.seq_length, self.hidden_size)				)





				def 				__UpperCAmelCase				( self							:	Union[str, Any]   , __lowerCamelCase							:	Any   , __lowerCamelCase							:	List[str]   , __lowerCamelCase							:	Optional[int]   , __lowerCamelCase							:	Dict   , __lowerCamelCase							:	Optional[int]   , __lowerCamelCase							:	int   , __lowerCamelCase							:	Union[str, Any]   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	Any   , )   ->	Optional[Any]:
							a       =     TFXLMWithLMHeadModel(__lowerCamelCase				)

							a       =     {"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
							a       =     model(__lowerCamelCase				)

							a       =     outputs

							self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.seq_length, self.vocab_size)				)





				def 				__UpperCAmelCase				( self							:	Tuple   , __lowerCamelCase							:	Dict   , __lowerCamelCase							:	Union[str, Any]   , __lowerCamelCase							:	int   , __lowerCamelCase							:	Union[str, Any]   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	Union[str, Any]   , __lowerCamelCase							:	Tuple   , __lowerCamelCase							:	Any   , __lowerCamelCase							:	List[Any]   , )   ->	int:
							a       =     TFXLMForQuestionAnsweringSimple(__lowerCamelCase				)

							a       =     {"input_ids": input_ids, "lengths": input_lengths}

							a       =     model(__lowerCamelCase				)

							self.parent.assertEqual(result.start_logits.shape   , (self.batch_size, self.seq_length)				)
							self.parent.assertEqual(result.end_logits.shape   , (self.batch_size, self.seq_length)				)





				def 				__UpperCAmelCase				( self							:	Dict   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	int   , __lowerCamelCase							:	Union[str, Any]   , __lowerCamelCase							:	List[str]   , __lowerCamelCase							:	List[Any]   , __lowerCamelCase							:	Tuple   , __lowerCamelCase							:	str   , __lowerCamelCase							:	Optional[int]   , __lowerCamelCase							:	Optional[int]   , )   ->	Optional[int]:
							a       =     TFXLMForSequenceClassification(__lowerCamelCase				)

							a       =     {"input_ids": input_ids, "lengths": input_lengths}

							a       =     model(__lowerCamelCase				)

							self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.type_sequence_label_size)				)





				def 				__UpperCAmelCase				( self							:	Any   , __lowerCamelCase							:	Tuple   , __lowerCamelCase							:	str   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	str   , __lowerCamelCase							:	Optional[int]   , __lowerCamelCase							:	List[str]   , __lowerCamelCase							:	Union[str, Any]   , __lowerCamelCase							:	List[str]   , __lowerCamelCase							:	Union[str, Any]   , )   ->	Optional[Any]:
							a       =     self.num_labels
							a       =     TFXLMForTokenClassification(config=__lowerCamelCase				)
							a       =     {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							a       =     model(__lowerCamelCase				)
							self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.seq_length, self.num_labels)				)





				def 				__UpperCAmelCase				( self							:	Any   , __lowerCamelCase							:	List[str]   , __lowerCamelCase							:	Union[str, Any]   , __lowerCamelCase							:	Dict   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	str   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	str   , __lowerCamelCase							:	List[str]   , __lowerCamelCase							:	Dict   , )   ->	List[Any]:
							a       =     self.num_choices
							a       =     TFXLMForMultipleChoice(config=__lowerCamelCase				)
							a       =     tf.tile(tf.expand_dims(__lowerCamelCase   , 1				)   , (1, self.num_choices, 1)				)
							a       =     tf.tile(tf.expand_dims(__lowerCamelCase   , 1				)   , (1, self.num_choices, 1)				)
							a       =     tf.tile(tf.expand_dims(__lowerCamelCase   , 1				)   , (1, self.num_choices, 1)				)
							a       =     {
							    "input_ids": multiple_choice_inputs_ids,
							    "attention_mask": multiple_choice_input_mask,
							    "token_type_ids": multiple_choice_token_type_ids,
							}
							a       =     model(__lowerCamelCase				)
							self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.num_choices)				)





				def 				__UpperCAmelCase				( self							:	Optional[int]				)   ->	str:
							a       =     self.prepare_config_and_inputs()
							(
							    (
							    a
							)       ,						(
							    a
							)       ,						(
							    a
							)       ,						(
							    a
							)       ,						(
							    a
							)       ,						(
							    a
							)       ,						(
							    a
							)       ,						(
							    a
							)       ,						(
							    a
							)       ,						
							)       =     config_and_inputs
							a       =     {
							    "input_ids": input_ids,
							    "token_type_ids": token_type_ids,
							    "langs": token_type_ids,
							    "lengths": input_lengths,
							}
							return config, inputs_dict




@require_tf
class        snake_case__				(_UpperCamelCase       ,			_UpperCamelCase       ,			unittest.TestCase  ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_			:				int									=   (
				    (
				        TFXLMModel,
				        TFXLMWithLMHeadModel,
				        TFXLMForSequenceClassification,
				        TFXLMForQuestionAnsweringSimple,
				        TFXLMForTokenClassification,
				        TFXLMForMultipleChoice,
				    )
				    if is_tf_available()
				    else ()
				)
				SCREAMING_SNAKE_CASE_			:				List[str]									=   (
				    (TFXLMWithLMHeadModel,) if is_tf_available() else ()
				)  # TODO (PVP): Check other models whether language generation is also applicable
				SCREAMING_SNAKE_CASE_			:				List[Any]									=   (
				    {
				        """feature-extraction""": TFXLMModel,
				        """fill-mask""": TFXLMWithLMHeadModel,
				        """question-answering""": TFXLMForQuestionAnsweringSimple,
				        """text-classification""": TFXLMForSequenceClassification,
				        """text-generation""": TFXLMWithLMHeadModel,
				        """token-classification""": TFXLMForTokenClassification,
				        """zero-shot""": TFXLMForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				SCREAMING_SNAKE_CASE_			:				Tuple									=   False
				SCREAMING_SNAKE_CASE_			:				Dict									=   False





				def 				__UpperCAmelCase				( self							:	Any   , __lowerCamelCase							:	Tuple   , __lowerCamelCase							:	List[Any]   , __lowerCamelCase							:	Any   , __lowerCamelCase							:	Optional[int]   , __lowerCamelCase							:	List[Any]				)   ->	Dict:
							if (
							    pipeline_test_casse_name == "QAPipelineTests"
							    and tokenizer_name is not None
							    and not tokenizer_name.endswith("Fast"				)
							):
										# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
										# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
										# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
										return True

							return False





				def 				__UpperCAmelCase				( self							:	Tuple				)   ->	Tuple:
							a       =     TFXLMModelTester(self				)
							a       =     ConfigTester(self   , config_class=__lowerCamelCase   , emb_dim=37				)





				def 				__UpperCAmelCase				( self							:	List[str]				)   ->	Optional[int]:
							self.config_tester.run_common_tests()





				def 				__UpperCAmelCase				( self							:	Dict				)   ->	List[Any]:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_xlm_model(*__lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	Tuple				)   ->	Dict:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_xlm_lm_head(*__lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	Optional[Any]				)   ->	Dict:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_xlm_qa(*__lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	List[Any]				)   ->	Optional[int]:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_xlm_sequence_classif(*__lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	Optional[int]				)   ->	Optional[Any]:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_xlm_for_token_classification(*__lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	List[Any]				)   ->	str:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_xlm_for_multiple_choice(*__lowerCamelCase				)





				@slow
				def 				__UpperCAmelCase				( self							:	Optional[Any]				)   ->	List[Any]:
							for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										a       =     TFXLMModel.from_pretrained(__lowerCamelCase				)
										self.assertIsNotNone(__lowerCamelCase				)




@require_tf
class        snake_case__				(unittest.TestCase  ):
				"""simple docstring"""



				@slow
				def 				__UpperCAmelCase				( self							:	Optional[Any]				)   ->	str:
							a       =     TFXLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048"				)
							a       =     tf.convert_to_tensor([[14, 4_47]]   , dtype=tf.intaa				)  # the president
							a       =     [
							    14,
							    4_47,
							    14,
							    4_47,
							    14,
							    4_47,
							    14,
							    4_47,
							    14,
							    4_47,
							    14,
							    4_47,
							    14,
							    4_47,
							    14,
							    4_47,
							    14,
							    4_47,
							    14,
							    4_47,
							]  # the president the president the president the president the president the president the president the president the president the president
							# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
							a       =     model.generate(__lowerCamelCase   , do_sample=__lowerCamelCase				)
							self.assertListEqual(output_ids[0].numpy().tolist()   , __lowerCamelCase				)






