import unittest

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

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 (
		    ErnieMForInformationExtraction,
		    ErnieMForMultipleChoice,
		    ErnieMForQuestionAnswering,
		    ErnieMForSequenceClassification,
		    ErnieMForTokenClassification,
		    ErnieMModel,
		)
		from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST




class        _UpperCamelCase				:


			"""simple docstring"""






			def __init__(       self				,       lowerCAmelCase__				,       lowerCAmelCase__=13				,       lowerCAmelCase__=7				,       lowerCAmelCase__=True				,       lowerCAmelCase__=True				,       lowerCAmelCase__=True				,       lowerCAmelCase__=99				,       lowerCAmelCase__=32				,       lowerCAmelCase__=5				,       lowerCAmelCase__=4				,       lowerCAmelCase__=37				,       lowerCAmelCase__="gelu"				,       lowerCAmelCase__=0.1				,       lowerCAmelCase__=0.1				,       lowerCAmelCase__=5_12				,       lowerCAmelCase__=16				,       lowerCAmelCase__=2				,       lowerCAmelCase__=0.02				,       lowerCAmelCase__=3				,       lowerCAmelCase__=4				,       lowerCAmelCase__=None				,       )					-> List[Any]:

							'''simple docstring'''


							__lowercase			      =							parent
							__lowercase			      =							batch_size
							__lowercase			      =							seq_length
							__lowercase			      =							is_training
							__lowercase			      =							use_input_mask
							__lowercase			      =							use_labels
							__lowercase			      =							vocab_size
							__lowercase			      =							hidden_size
							__lowercase			      =							num_hidden_layers
							__lowercase			      =							num_attention_heads
							__lowercase			      =							intermediate_size
							__lowercase			      =							hidden_act
							__lowercase			      =							hidden_dropout_prob
							__lowercase			      =							attention_probs_dropout_prob
							__lowercase			      =							max_position_embeddings
							__lowercase			      =							type_vocab_size
							__lowercase			      =							type_sequence_label_size
							__lowercase			      =							initializer_range
							__lowercase			      =							num_labels
							__lowercase			      =							num_choices
							__lowercase			      =							scope





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> str:

							'''simple docstring'''


							__lowercase			      =							ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size						)

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

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

							__lowercase			      =							self.get_config()

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





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Tuple:

							'''simple docstring'''


							__lowercase			      =							ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size						)

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

							return config, input_ids, input_mask





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Dict:

							'''simple docstring'''


							return ErnieMConfig(
							    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				,       initializer_range=self.initializer_range				,       )





			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__						)					-> Any:

							'''simple docstring'''


							__lowercase			      =							ErnieMModel(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase			      =							model(lowerCAmelCase__				,       return_dict=lowerCAmelCase__						)
							self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)						)





			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__						)					-> int:

							'''simple docstring'''


							__lowercase			      =							ErnieMForQuestionAnswering(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase			      =							model(
							    lowerCAmelCase__				,       attention_mask=lowerCAmelCase__				,       start_positions=lowerCAmelCase__				,       end_positions=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 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__						)					-> List[str]:

							'''simple docstring'''


							__lowercase			      =							ErnieMForInformationExtraction(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase			      =							torch.ones_like(lowerCAmelCase__				,       dtype=torch.floataa						)
							__lowercase			      =							model(
							    lowerCAmelCase__				,       attention_mask=lowerCAmelCase__				,       start_positions=lowerCAmelCase__				,       end_positions=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 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__						)					-> List[str]:

							'''simple docstring'''


							__lowercase			      =							self.num_labels
							__lowercase			      =							ErnieMForSequenceClassification(lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase			      =							model(lowerCAmelCase__				,       attention_mask=lowerCAmelCase__				,       labels=lowerCAmelCase__						)
							self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)						)





			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__						)					-> Union[str, Any]:

							'''simple docstring'''


							__lowercase			      =							self.num_labels
							__lowercase			      =							ErnieMForTokenClassification(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							input_ids.to(lowerCAmelCase__						)
							input_mask.to(lowerCAmelCase__						)
							token_labels.to(lowerCAmelCase__						)

							__lowercase			      =							model(lowerCAmelCase__				,       attention_mask=lowerCAmelCase__				,       labels=lowerCAmelCase__						)

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





			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__						)					-> List[Any]:

							'''simple docstring'''


							__lowercase			      =							self.num_choices
							__lowercase			      =							ErnieMForMultipleChoice(config=lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							__lowercase			      =							input_ids.unsqueeze(1						).expand(-1				,       self.num_choices				,       -1						).contiguous()
							__lowercase			      =							input_mask.unsqueeze(1						).expand(-1				,       self.num_choices				,       -1						).contiguous()
							__lowercase			      =							model(
							    lowerCAmelCase__				,       attention_mask=lowerCAmelCase__				,       labels=lowerCAmelCase__				,       )
							self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_choices)						)





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Optional[int]:

							'''simple docstring'''


							__lowercase			      =							self.prepare_config_and_inputs()
							(
							    (
							    __lowercase
							)      ,					(
							    __lowercase
							)      ,					(
							    __lowercase
							)      ,					(
							    __lowercase
							)      ,					(
							    __lowercase
							)      ,					(
							    __lowercase
							)      ,					
							)			      =							config_and_inputs
							__lowercase			      =							{'''input_ids''': input_ids, '''attention_mask''': input_mask}
							return config, inputs_dict




@require_torch
class        _UpperCamelCase				(						_UpperCAmelCase						,_UpperCAmelCase						,unittest.TestCase			):


			"""simple docstring"""






			__a  :			Optional[int]      			=						(
			    (
			        ErnieMModel,
			        ErnieMForMultipleChoice,
			        ErnieMForQuestionAnswering,
			        ErnieMForSequenceClassification,
			        ErnieMForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			__a  :			int      			=						()
			__a  :			Dict      			=						(
			    {
			        '''feature-extraction''': ErnieMModel,
			        '''question-answering''': ErnieMForQuestionAnswering,
			        '''text-classification''': ErnieMForSequenceClassification,
			        '''token-classification''': ErnieMForTokenClassification,
			        '''zero-shot''': ErnieMForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			__a  :			str      			=						False





			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__				,       lowerCAmelCase__						)					-> str:

							'''simple docstring'''


							if pipeline_test_casse_name == "QAPipelineTests":
											return True

							return False





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Tuple:

							'''simple docstring'''


							__lowercase			      =							ErnieMModelTester(self						)
							__lowercase			      =							ConfigTester(self				,       config_class=lowerCAmelCase__				,       hidden_size=37						)





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> List[Any]:

							'''simple docstring'''


							self.config_tester.run_common_tests()





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Optional[Any]:

							'''simple docstring'''


							__lowercase			      =							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowerCAmelCase__						)





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Any:

							'''simple docstring'''


							__lowercase			      =							self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
											__lowercase			      =							type
											self.model_tester.create_and_check_model(*lowerCAmelCase__						)





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Optional[Any]:

							'''simple docstring'''


							__lowercase			      =							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase__						)





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Dict:

							'''simple docstring'''


							__lowercase			      =							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase__						)





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Optional[int]:

							'''simple docstring'''


							__lowercase			      =							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase__						)





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Dict:

							'''simple docstring'''


							__lowercase			      =							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_information_extraction(*lowerCAmelCase__						)





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Tuple:

							'''simple docstring'''


							__lowercase			      =							self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__						)





			@slow
			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Tuple:

							'''simple docstring'''


							for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__lowercase			      =							ErnieMModel.from_pretrained(lowerCAmelCase__						)
											self.assertIsNotNone(lowerCAmelCase__						)




@require_torch
class        _UpperCamelCase				(						unittest.TestCase			):


			"""simple docstring"""






			@slow
			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Tuple:

							'''simple docstring'''


							__lowercase			      =							ErnieMModel.from_pretrained('''susnato/ernie-m-base_pytorch'''						)
							model.eval()
							__lowercase			      =							torch.tensor([[0, 1, 2, 3, 4, 5]]						)
							__lowercase			      =							model(lowerCAmelCase__						)[0]

							# TODO Replace vocab size
							__lowercase			      =							7_68

							__lowercase			      =							torch.Size((1, 6, hidden_size)						)
							self.assertEqual(output.shape				,       lowerCAmelCase__						)

							__lowercase			      =							torch.tensor(
							    [[[-0.0012, 0.1245, -0.0214], [-0.0742, 0.0244, -0.0771], [-0.0333, 0.1164, -0.1554]]]						)

							self.assertTrue(torch.allclose(output[:, :3, :3]				,       lowerCAmelCase__				,       atol=1E-3						)						)
import jax
import jax.numpy as jnp
from bigbird_flax import FlaxBigBirdForNaturalQuestions
from datasets import load_from_disk

from transformers import BigBirdTokenizerFast


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





def     UpperCAmelCase						(   lowercase							,      lowercase=0							,      lowercase=None							):







				"""simple docstring"""



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





def     UpperCAmelCase						(   lowercase							,      lowercase=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			      =							(
								    given_answers + get_sub_answers(lowercase							,      begin=1							) + get_sub_answers(lowercase							,      end=-1							)
								)
				__lowercase			      =							[]
				for answer in given_answers:
								__lowercase			      =							answer.replace('''_'''							,      ''' '''							).lower()
								__lowercase			      =							''''''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ''' ''' for c in alias							)
								answers.append(''' '''.join(alias.split()							).strip()							)
				return set(lowercase							)





def     UpperCAmelCase						(   lowercase							,      lowercase							,      lowercase=1							,      lowercase=100							):







				"""simple docstring"""



				__lowercase      ,					__lowercase			      =							jax.lax.top_k(lowercase							,      lowercase							)
				__lowercase      ,					__lowercase			      =							jax.lax.top_k(lowercase							,      lowercase							)

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

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





def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							sample['''question''']['''text''']
				__lowercase			      =							sample['''document''']['''tokens''']['''token''']
				__lowercase			      =							sample['''document''']['''tokens''']['''is_html''']
				__lowercase			      =							sample['''annotations''']['''long_answer''']
				__lowercase			      =							sample['''annotations''']['''short_answers''']

				__lowercase			      =							''' '''.join([context[i] for i in range(len(lowercase							)							) 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			      =							[]
				for s in short_answers:
								short_targets.extend(s['''text''']							)
				__lowercase			      =							list(set(lowercase							)							)

				__lowercase			      =							[]
				for s in long_answers:
								if s["start_token"] == -1:
												continue
								__lowercase			      =							context[s['''start_token'''] : s['''end_token''']]
								__lowercase			      =							is_html[s['''start_token'''] : s['''end_token''']]
								__lowercase			      =							''' '''.join([answer[i] for i in range(len(lowercase							)							) if not html[i]]							)
								if new_answer not in long_targets:
												long_targets.append(lowercase							)

				__lowercase			      =							'''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     UpperCAmelCase						(   ):







				"""simple docstring"""



				__lowercase			      =							load_from_disk('''natural-questions-validation'''							)
				__lowercase			      =							dataset.map(lowercase							).remove_columns(['''annotations''', '''document''', '''id''']							)
				print(lowercase							)

				__lowercase			      =							dataset.filter(lambda lowercase							: (len(x['''question''']							) + len(x['''context''']							)) < 4 * 4096							)
				__lowercase			      =							short_validation_dataset.filter(lambda lowercase							: x["category"] != "null"							)
				short_validation_dataset

				__lowercase			      =							'''vasudevgupta/flax-bigbird-natural-questions'''
				__lowercase			      =							FlaxBigBirdForNaturalQuestions.from_pretrained(lowercase							)
				__lowercase			      =							BigBirdTokenizerFast.from_pretrained(lowercase							)

				@jax.jit
				def forward(*lowercase							,      **lowercase							):
								__lowercase      ,					__lowercase      ,					__lowercase			      =							model(*lowercase							,      **lowercase							)
								return start_logits, end_logits, jnp.argmax(lowercase							,      axis=-1							)

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

								__lowercase      ,					__lowercase      ,					__lowercase			      =							forward(**lowercase							)

								__lowercase			      =							CATEGORY_MAPPING[category.item()]

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

								if predicted_category in ["yes", "no", "null"]:
												__lowercase			      =							[predicted_category]
												__lowercase			      =							example['''output'''] == example['''targets''']
												__lowercase			      =							predicted_category != '''null'''
												return example

								__lowercase			      =							38 if predicted_category == '''short''' else 1024
								__lowercase      ,					__lowercase			      =							get_best_valid_start_end_idx(
								    start_scores[0]							,      end_scores[0]							,      top_k=8							,      max_size=lowercase							)

								__lowercase			      =							inputs['''input_ids'''][0].tolist()
								__lowercase			      =							[tokenizer.decode(input_ids[start_score : end_score + 1]							)]

								__lowercase			      =							expand_to_aliases(example['''targets''']							,      make_sub_answers=lowercase							)
								__lowercase			      =							expand_to_aliases(example['''output''']							)

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

								# if there is a common element, it's a exact match
								__lowercase			      =							len(list(answers & predictions							)							) > 0
								__lowercase			      =							predicted_category != '''null''' and len(lowercase							) > 0

								return example

				__lowercase			      =							short_validation_dataset.map(lowercase							)

				__lowercase			      =							len(lowercase							)
				__lowercase			      =							len(short_validation_dataset.filter(lambda lowercase							: x["match"] == 1							)							)
				print('''EM score:'''							,      (matched / total) * 100							,      '''%'''							)


if __name__ == "__main__":
		main()