

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

from transformers import BigBirdTokenizerFast


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






def        lowerCamelCase_  (				UpperCAmelCase_				:  Any				,							UpperCAmelCase_				:  int=0				,							UpperCAmelCase_				:  List[str]=None						)    ->					Optional[int]:




	'''simple docstring'''




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






def        lowerCamelCase_  (				UpperCAmelCase_				:  Union[str, Any]				,							UpperCAmelCase_				:  Tuple=False						)    ->					Tuple:




	'''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"
		_UpperCamelCase :		int								=    (
		    given_answers + get_sub_answers(UpperCAmelCase_				,							begin=1						) + get_sub_answers(UpperCAmelCase_				,							end=-1						)
		)
	_UpperCamelCase :		Union[str, Any]								=    []
	for answer in given_answers:
		_UpperCamelCase :		str								=    answer.replace('_'				,							' '						).lower()
		_UpperCamelCase :		Any								=    ''.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_				:  str				,							UpperCAmelCase_				:  Union[str, Any]				,							UpperCAmelCase_				:  List[Any]=1				,							UpperCAmelCase_				:  Optional[int]=1_0_0						)    ->					str:




	'''simple docstring'''




	_UpperCamelCase						,		_UpperCamelCase :		Any								=    jax.lax.top_k(UpperCAmelCase_				,							UpperCAmelCase_						)
	_UpperCamelCase						,		_UpperCamelCase :		Union[str, Any]								=    jax.lax.top_k(UpperCAmelCase_				,							UpperCAmelCase_						)

	_UpperCamelCase :		int								=    best_end_idx[:, None] - best_start_idx[None, :]
	_UpperCamelCase :		Optional[Any]								=    jnp.logical_or(widths < 0				,							widths > max_size						)
	_UpperCamelCase :		List[str]								=    (best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
	_UpperCamelCase :		int								=    jnp.argmax(UpperCAmelCase_						).item()

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






def        lowerCamelCase_  (				UpperCAmelCase_				:  List[str]						)    ->					Any:




	'''simple docstring'''




	_UpperCamelCase :		Optional[Any]								=    sample['question']['text']
	_UpperCamelCase :		List[str]								=    sample['document']['tokens']['token']
	_UpperCamelCase :		List[Any]								=    sample['document']['tokens']['is_html']
	_UpperCamelCase :		List[Any]								=    sample['annotations']['long_answer']
	_UpperCamelCase :		Dict								=    sample['annotations']['short_answers']

	_UpperCamelCase :		Union[str, Any]								=    ' '.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",
			}

	_UpperCamelCase :		int								=    []
	for s in short_answers:
		short_targets.extend(s['text']						)
	_UpperCamelCase :		str								=    list(set(UpperCAmelCase_						)						)

	_UpperCamelCase :		Optional[int]								=    []
	for s in long_answers:
		if s["start_token"] == -1:
			continue
		_UpperCamelCase :		List[str]								=    context[s['start_token'] : s['end_token']]
		_UpperCamelCase :		List[str]								=    is_html[s['start_token'] : s['end_token']]
		_UpperCamelCase :		Union[str, 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_						)

	_UpperCamelCase :		Union[str, Any]								=    '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_  (				)    ->					Union[str, Any]:




	'''simple docstring'''




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

	_UpperCamelCase :		Union[str, Any]								=    dataset.filter(lambda UpperCAmelCase_						: (len(x['question']						) + len(x['context']						)) < 4 * 4_0_9_6						)
	_UpperCamelCase :		Optional[Any]								=    short_validation_dataset.filter(lambda UpperCAmelCase_						: x["category"] != "null"						)
	short_validation_dataset

	_UpperCamelCase :		int								=    'vasudevgupta/flax-bigbird-natural-questions'
	_UpperCamelCase :		List[str]								=    FlaxBigBirdForNaturalQuestions.from_pretrained(UpperCAmelCase_						)
	_UpperCamelCase :		str								=    BigBirdTokenizerFast.from_pretrained(UpperCAmelCase_						)

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

	def evaluate(UpperCAmelCase_				:  Union[str, Any]						):
		# encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
		_UpperCamelCase :		Dict								=    tokenizer(
		    example['question']				,							example['context']				,							return_tensors='np'				,							max_length=4_0_9_6				,							padding='max_length'				,							truncation=UpperCAmelCase_				,							)

		_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase :		Union[str, Any]								=    forward(**UpperCAmelCase_						)

		_UpperCamelCase :		Optional[Any]								=    CATEGORY_MAPPING[category.item()]

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

		if predicted_category in ["yes", "no", "null"]:
			_UpperCamelCase :		Dict								=    [predicted_category]
			_UpperCamelCase :		str								=    example['output'] == example['targets']
			_UpperCamelCase :		int								=    predicted_category != 'null'
			return example

		_UpperCamelCase :		Tuple								=    3_8 if predicted_category == 'short' else 1_0_2_4
		_UpperCamelCase						,		_UpperCamelCase :		List[Any]								=    get_best_valid_start_end_idx(
		    start_scores[0]				,							end_scores[0]				,							top_k=8				,							max_size=UpperCAmelCase_						)

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

		_UpperCamelCase :		Tuple								=    expand_to_aliases(example['targets']				,							make_sub_answers=UpperCAmelCase_						)
		_UpperCamelCase :		Tuple								=    expand_to_aliases(example['output']						)

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

		# if there is a common element, it's a exact match
		_UpperCamelCase :		Dict								=    len(list(answers & predictions						)						) > 0
		_UpperCamelCase :		Any								=    predicted_category != 'null' and len(UpperCAmelCase_						) > 0

		return example

	_UpperCamelCase :		Optional[int]								=    short_validation_dataset.map(UpperCAmelCase_						)

	_UpperCamelCase :		Tuple								=    len(UpperCAmelCase_						)
	_UpperCamelCase :		List[str]								=    len(short_validation_dataset.filter(lambda UpperCAmelCase_						: x["match"] == 1						)						)
	print('EM score:'				,							(matched / total) * 1_0_0				,							'%'						)


if __name__ == "__main__":
							main()







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowerCAmelCase__             =		{
    """configuration_markuplm""": ["""MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MarkupLMConfig"""],
    """feature_extraction_markuplm""": ["""MarkupLMFeatureExtractor"""],
    """processing_markuplm""": ["""MarkupLMProcessor"""],
    """tokenization_markuplm""": ["""MarkupLMTokenizer"""],
}

try:
							if not is_tokenizers_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCAmelCase__             =		["""MarkupLMTokenizerFast"""]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCAmelCase__             =		[
							    """MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """MarkupLMForQuestionAnswering""",
							    """MarkupLMForSequenceClassification""",
							    """MarkupLMForTokenClassification""",
							    """MarkupLMModel""",
							    """MarkupLMPreTrainedModel""",
							]


if TYPE_CHECKING:
							from .configuration_markuplm import MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP, MarkupLMConfig
							from .feature_extraction_markuplm import MarkupLMFeatureExtractor
							from .processing_markuplm import MarkupLMProcessor
							from .tokenization_markuplm import MarkupLMTokenizer

							try:
														if not is_tokenizers_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_markuplm_fast import MarkupLMTokenizerFast

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_markuplm import (
														    MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST,
														    MarkupLMForQuestionAnswering,
														    MarkupLMForSequenceClassification,
														    MarkupLMForTokenClassification,
														    MarkupLMModel,
														    MarkupLMPreTrainedModel,
														)


else:
							import sys

							lowerCAmelCase__             =		_LazyModule(__name__, globals()["""__file__"""], _import_structure)




