
'''simple docstring'''





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

from transformers import BigBirdTokenizerFast


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






def  a_	( __lowercase     :			List[str] ,      __lowercase     :			List[str]=0 ,      __lowercase     :			Union[str, Any]=None				)				->						List[Any]:
			return [" ".join(x.split(' '				)[begin:end]				) for x in answers if len(x.split(' '				)				) > 1]






def  a_	( __lowercase     :			int ,      __lowercase     :			Dict=False				)				->						Union[str, Any]:
			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"
						_snake_case						=    (
						    given_answers + get_sub_answers(lowercase__ ,      begin=1				) + get_sub_answers(lowercase__ ,      end=-1				)
						)
			_snake_case						=    []
			for answer in given_answers:
						_snake_case						=    answer.replace('_' ,      ' '				).lower()
						_snake_case						=    "".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  a_	( __lowercase     :			List[Any] ,      __lowercase     :			Tuple ,      __lowercase     :			List[Any]=1 ,      __lowercase     :			Any=100				)				->						List[Any]:
			_snake_case						=    jax.lax.top_k(lowercase__ ,      lowercase__				)
			_snake_case						=    jax.lax.top_k(lowercase__ ,      lowercase__				)

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

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






def  a_	( __lowercase     :			str				)				->						Union[str, Any]:
			_snake_case						=    sample["question"]["text"]
			_snake_case						=    sample["document"]["tokens"]["token"]
			_snake_case						=    sample["document"]["tokens"]["is_html"]
			_snake_case						=    sample["annotations"]["long_answer"]
			_snake_case						=    sample["annotations"]["short_answers"]

			_snake_case						=    " ".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",
									}

			_snake_case						=    []
			for s in short_answers:
						short_targets.extend(s['text']				)
			_snake_case						=    list(set(lowercase__				)				)

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

			_snake_case						=    "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  a_	( )				->						List[Any]:
			_snake_case						=    load_from_disk('natural-questions-validation'				)
			_snake_case						=    dataset.map(lowercase__				).remove_columns(['annotations', 'document', 'id']				)
			print(lowercase__				)

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

			_snake_case						=    "vasudevgupta/flax-bigbird-natural-questions"
			_snake_case						=    FlaxBigBirdForNaturalQuestions.from_pretrained(lowercase__				)
			_snake_case						=    BigBirdTokenizerFast.from_pretrained(lowercase__				)

			@jax.jit
			def forward(*__lowercase     :			Union[str, Any] ,      **__lowercase     :			int				):
						_snake_case						=    model(*lowercase__ ,      **lowercase__				)
						return start_logits, end_logits, jnp.argmax(lowercase__ ,      axis=-1				)

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

						_snake_case						=    forward(**lowercase__				)

						_snake_case						=    CATEGORY_MAPPING[category.item()]

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

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

						_snake_case						=    38 if predicted_category == "short" else 1_024
						_snake_case						=    get_best_valid_start_end_idx(
						    start_scores[0] ,      end_scores[0] ,      top_k=8 ,      max_size=lowercase__				)

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

						_snake_case						=    expand_to_aliases(example['targets'] ,      make_sub_answers=lowercase__				)
						_snake_case						=    expand_to_aliases(example['output']				)

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

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

						return example

			_snake_case						=    short_validation_dataset.map(lowercase__				)

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


if __name__ == "__main__":
			main()

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding




class      SCREAMING_SNAKE_CASE__  (			UpperCAmelCase      ):





	'''simple docstring'''



	_UpperCAmelCase :	int					     =      ["image_processor", "tokenizer"]
	_UpperCAmelCase :	Tuple					     =      "EfficientNetImageProcessor"
	_UpperCAmelCase :	Optional[Any]					     =      ("BertTokenizer", "BertTokenizerFast")



	def __init__(						self			:					Tuple    ,						lowercase			:					List[Any]    ,						lowercase			:					Optional[int]			):


				'''simple docstring'''


				super().__init__(lowercase    ,						lowercase			)



	def __call__(						self			:					Union[str, Any]    ,						lowercase			:					Dict=None    ,						lowercase			:					Optional[int]=None    ,						lowercase			:					List[Any]="max_length"    ,						lowercase			:					Optional[int]=64    ,						lowercase			:					int=None    ,						**lowercase			:					Tuple			):


				'''simple docstring'''


				if text is None and images is None:
							raise ValueError('You have to specify either text or images. Both cannot be none.'			)

				if text is not None:
							_snake_case						=    self.tokenizer(
							    lowercase    ,						padding=lowercase    ,						max_length=lowercase    ,						return_tensors=lowercase    ,						**lowercase			)

				if images is not None:
							_snake_case						=    self.image_processor(lowercase    ,						return_tensors=lowercase    ,						**lowercase			)

				if text is not None and images is not None:
							_snake_case						=    image_features.pixel_values
							return encoding
				elif text is not None:
							return encoding
				else:
							return BatchEncoding(data=dict(**lowercase			)    ,						tensor_type=lowercase			)



	def 				A							(						self			:					Any    ,						*lowercase			:					List[str]    ,						**lowercase			:					Optional[Any]			):


				'''simple docstring'''


				return self.tokenizer.batch_decode(*lowercase    ,						**lowercase			)



	def 				A							(						self			:					List[Any]    ,						*lowercase			:					Union[str, Any]    ,						**lowercase			:					Tuple			):


				'''simple docstring'''


				return self.tokenizer.decode(*lowercase    ,						**lowercase			)





	@property
	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


				_snake_case						=    self.tokenizer.model_input_names
				_snake_case						=    self.image_processor.model_input_names
				return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names			)			)