
'''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  							= {0: 'null', 1: 'short', 2: 'long', 3: 'yes', 4: 'no'}
__lowerCAmelCase  							= set(''.join(['‘', '’', '´', '`', '.', ',', '-', '"']))






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE=0							,      _SCREAMING_SNAKE_CASE=None   ):
							return [" ".join(x.split(""" """   )[begin:end]   ) for x in answers if len(x.split(""" """   )   ) > 1]






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE=False   ):
							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(_SCREAMING_SNAKE_CASE							,      begin=1   ) + get_sub_answers(_SCREAMING_SNAKE_CASE							,      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(_SCREAMING_SNAKE_CASE   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE=1							,      _SCREAMING_SNAKE_CASE=100   ):
							_snake_case,							_snake_case     =						jax.lax.top_k(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   )
							_snake_case,							_snake_case     =						jax.lax.top_k(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   )

							_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(_SCREAMING_SNAKE_CASE   ).item()

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






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE   ):
							_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(_SCREAMING_SNAKE_CASE   )   ) 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(_SCREAMING_SNAKE_CASE   )   )

							_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(_SCREAMING_SNAKE_CASE   )   ) if not html[i]]   )
														if new_answer not in long_targets:
																					long_targets.append(_SCREAMING_SNAKE_CASE   )

							_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       __SCREAMING_SNAKE_CASE   (						):
							_snake_case     =						load_from_disk("""natural-questions-validation"""   )
							_snake_case     =						dataset.map(_SCREAMING_SNAKE_CASE   ).remove_columns(["""annotations""", """document""", """id"""]   )
							print(_SCREAMING_SNAKE_CASE   )

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

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

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

							def evaluate(_SCREAMING_SNAKE_CASE   ):
														# 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=4096							,      padding="""max_length"""							,      truncation=_SCREAMING_SNAKE_CASE							,      )

														_snake_case,							_snake_case,							_snake_case     =						forward(**_SCREAMING_SNAKE_CASE   )

														_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 1024
														_snake_case,							_snake_case     =						get_best_valid_start_end_idx(
														    start_scores[0]							,      end_scores[0]							,      top_k=8							,      max_size=_SCREAMING_SNAKE_CASE   )

														_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=_SCREAMING_SNAKE_CASE   )
														_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(_SCREAMING_SNAKE_CASE   ) > 0

														return example

							_snake_case     =						short_validation_dataset.map(_SCREAMING_SNAKE_CASE   )

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


if __name__ == "__main__":
				main()

'''simple docstring'''
import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


__lowerCAmelCase  							= logging.get_logger(__name__)


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def __init__(self						)  ->     Any:
														super().__init__()
														if version.parse(torch.__version__						) < version.parse("""1.12.0"""						):
																					raise ImportError(
																					    f"""You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use """
																					    """PytorchGELUTanh. Please upgrade torch."""						)





							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return nn.functional.gelu(UpperCAmelCase						,						approximate="""tanh"""						)


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi						) * (input + 0.04_4715 * torch.pow(UpperCAmelCase						,						3.0						))						))


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def __init__(self						,						UpperCAmelCase = False						)  ->     Optional[Any]:
														super().__init__()
														if use_gelu_python:
																					_snake_case     =						self._gelu_python
														else:
																					_snake_case     =						nn.functional.gelu

							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0						)						))





							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return self.act(UpperCAmelCase						)


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return 0.5 * input * (1.0 + torch.tanh(input * 0.79_7884_5608 * (1.0 + 0.04_4715 * input * input)						))


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return input * torch.sigmoid(1.702 * input						)


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def __init__(self						,						UpperCAmelCase						,						UpperCAmelCase						)  ->     List[Any]:
														if min > max:
																					raise ValueError(f"""min should be < max (got min: {min}, max: {max})"""						)

														super().__init__()
														_snake_case     =						min
														_snake_case     =						max





							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return torch.clip(gelu(UpperCAmelCase						)						,						self.min						,						self.max						)


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def __init__(self						)  ->     Tuple:
														super().__init__()
														_snake_case     =						math.sqrt(2 / math.pi						)





							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.04_4715 * torch.pow(UpperCAmelCase						,						3						))						))


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return nn.functional.silu(UpperCAmelCase						)


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def __init__(self						)  ->     Any:
														super().__init__()
														if version.parse(torch.__version__						) < version.parse("""1.9.0"""						):
																					_snake_case     =						self._mish_python
														else:
																					_snake_case     =						nn.functional.mish

							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return input * torch.tanh(nn.functional.softplus(UpperCAmelCase						)						)





							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return self.act(UpperCAmelCase						)


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def  lowercase							(self						,						UpperCAmelCase						)  ->     Tensor:
														return input


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def  lowercase							(self						,						UpperCAmelCase						,						UpperCAmelCase=0.70_7107						,						UpperCAmelCase=0.28_2095						)  ->     List[str]:
														_snake_case     =						(input - mu).div(sigma * math.sqrt(2.0						)						)
														return 0.5 * (1.0 + torch.erf(UpperCAmelCase						))


class     _lowerCAmelCase					(					nn.Module   ):







							'''simple docstring'''

							def  lowercase							(self						,						UpperCAmelCase						)  ->     Any:
														_snake_case     =						nn.functional.relu(UpperCAmelCase						)
														_snake_case     =						torch.square(UpperCAmelCase						)
														return squared


class     _lowerCAmelCase					(					__snake_case   ):







							'''simple docstring'''

							def __getitem__(self						,						UpperCAmelCase						)  ->     Tuple:
														_snake_case     =						super().__getitem__(UpperCAmelCase						)
														_snake_case,							_snake_case     =						content if isinstance(UpperCAmelCase						,						UpperCAmelCase						) else (content, {})
														return cls(**UpperCAmelCase						)
__lowerCAmelCase  							= {
    'gelu': GELUActivation,
    'gelu_10': (ClippedGELUActivation, {'min': -10, 'max': 10}),
    'gelu_fast': FastGELUActivation,
    'gelu_new': NewGELUActivation,
    'gelu_python': (GELUActivation, {'use_gelu_python': True}),
    'gelu_pytorch_tanh': PytorchGELUTanh,
    'gelu_accurate': AccurateGELUActivation,
    'laplace': LaplaceActivation,
    'linear': LinearActivation,
    'mish': MishActivation,
    'quick_gelu': QuickGELUActivation,
    'relu': nn.ReLU,
    'relu2': ReLUSquaredActivation,
    'relu6': nn.ReLUa,
    'sigmoid': nn.Sigmoid,
    'silu': SiLUActivation,
    'swish': SiLUActivation,
    'tanh': nn.Tanh,
}
__lowerCAmelCase  							= ClassInstantier(ACTaCLS)






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE   ):
							if activation_string in ACTaFN:
														return ACTaFN[activation_string]
							else:
														raise KeyError(f"""function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()   )}"""   )


# For backwards compatibility with: from activations import gelu_python
__lowerCAmelCase  							= get_activation('gelu_python')
__lowerCAmelCase  							= get_activation('gelu_new')
__lowerCAmelCase  							= get_activation('gelu')
__lowerCAmelCase  							= get_activation('gelu_fast')
__lowerCAmelCase  							= get_activation('quick_gelu')
__lowerCAmelCase  							= get_activation('silu')
__lowerCAmelCase  							= get_activation('mish')

__lowerCAmelCase  							= get_activation('linear')