



"""simple docstring"""




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

from transformers import BigBirdTokenizerFast


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



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



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

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

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



def  lowercase    (				SCREAMING_SNAKE_CASE__    :		Any			)					->       Dict:
 _snake_case      :					int					=   sample["question"]["text"]
 _snake_case      :					List[Any]					=   sample["document"]["tokens"]["token"]
 _snake_case      :					Optional[Any]					=   sample["document"]["tokens"]["is_html"]
 _snake_case      :					Optional[Any]					=   sample["annotations"]["long_answer"]
 _snake_case      :					str					=   sample["annotations"]["short_answers"]

 _snake_case      :					Dict					=   " ".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      :					Dict					=   []
 for s in short_answers:
  short_targets.extend(s["""text"""]			)
 _snake_case      :					List[str]					=   list(set(SCREAMING_SNAKE_CASE__			)			)

 _snake_case      :					Any					=   []
 for s in long_answers:
  if s["start_token"] == -1:
   continue
  _snake_case      :					Tuple					=   context[s["start_token"] : s["end_token"]]
  _snake_case      :					Tuple					=   is_html[s["start_token"] : s["end_token"]]
  _snake_case      :					Optional[int]					=   " ".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      :					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  lowercase    (				)					->       Union[str, Any]:
 _snake_case      :					Dict					=   load_from_disk("""natural-questions-validation"""			)
 _snake_case      :					Optional[Any]					=   dataset.map(SCREAMING_SNAKE_CASE__			).remove_columns(["""annotations""", """document""", """id"""]			)
 print(SCREAMING_SNAKE_CASE__			)

 _snake_case      :					Union[str, Any]					=   dataset.filter(lambda SCREAMING_SNAKE_CASE__			: (len(x["""question"""]			) + len(x["""context"""]			)) < 4 * 4_096			)
 _snake_case      :					Any					=   short_validation_dataset.filter(lambda SCREAMING_SNAKE_CASE__			: x["category"] != "null"			)
 short_validation_dataset

 _snake_case      :					str					=   "vasudevgupta/flax-bigbird-natural-questions"
 _snake_case      :					List[str]					=   FlaxBigBirdForNaturalQuestions.from_pretrained(SCREAMING_SNAKE_CASE__			)
 _snake_case      :					Dict					=   BigBirdTokenizerFast.from_pretrained(SCREAMING_SNAKE_CASE__			)

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

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

  _snake_case      :					List[Any]					=   forward(**SCREAMING_SNAKE_CASE__			)

  _snake_case      :					str					=   CATEGORY_MAPPING[category.item()]

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

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

  _snake_case      :					List[Any]					=   38 if predicted_category == "short" else 1_024
  _snake_case      :					Tuple					=   get_best_valid_start_end_idx(
      start_scores[0]			,    end_scores[0]			,    top_k=8			,    max_size=SCREAMING_SNAKE_CASE__			)

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

  _snake_case      :					Optional[int]					=   expand_to_aliases(example["""targets"""]			,    make_sub_answers=SCREAMING_SNAKE_CASE__			)
  _snake_case      :					Union[str, Any]					=   expand_to_aliases(example["""output"""]			)

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

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

  return example

 _snake_case      :					List[str]					=   short_validation_dataset.map(SCREAMING_SNAKE_CASE__			)

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


if __name__ == "__main__":
  main()





import math
from collections import OrderedDict

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

from .utils import logging


a__      =		logging.get_logger(__name__)







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def __init__(   self							:     Tuple)			->			Optional[int]:

    """simple docstring"""







    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        UpperCamelCase_					(   self							:     Tuple			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return nn.functional.gelu(lowerCAmelCase			,						approximate="""tanh""")







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def        UpperCamelCase_					(   self							:     Optional[int]			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044_715 * torch.pow(lowerCAmelCase			,						3.0))))







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def __init__(   self							:     str			,						lowerCAmelCase							:     bool = False)			->			Optional[Any]:

    """simple docstring"""







    super().__init__()
    if use_gelu_python:
     _snake_case      :					Optional[int]					=   self._gelu_python
    else:
     _snake_case      :					Any					=   nn.functional.gelu

   def        UpperCamelCase_					(   self							:     Union[str, Any]			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0)))


   def        UpperCamelCase_					(   self							:     str			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return self.act(lowerCAmelCase)







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def        UpperCamelCase_					(   self							:     int			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return 0.5 * input * (1.0 + torch.tanh(input * 0.7_978_845_608 * (1.0 + 0.044_715 * input * input)))







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def        UpperCamelCase_					(   self							:     int			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return input * torch.sigmoid(1.702 * input)







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def __init__(   self							:     int			,						lowerCAmelCase							:     float			,						lowerCAmelCase							:     float)			->			Optional[int]:

    """simple docstring"""







    if min > max:
     raise ValueError(F'''min should be < max (got min: {min}, max: {max})''')

    super().__init__()
    _snake_case      :					Dict					=   min
    _snake_case      :					Optional[int]					=   max


   def        UpperCamelCase_					(   self							:     Union[str, Any]			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return torch.clip(gelu(lowerCAmelCase)			,						self.min			,						self.max)







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def __init__(   self							:     Tuple)			->			List[str]:

    """simple docstring"""







    super().__init__()
    _snake_case      :					Union[str, Any]					=   math.sqrt(2 / math.pi)


   def        UpperCamelCase_					(   self							:     Dict			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.044_715 * torch.pow(lowerCAmelCase			,						3))))







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def        UpperCamelCase_					(   self							:     Dict			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return nn.functional.silu(lowerCAmelCase)







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def __init__(   self							:     List[Any])			->			str:

    """simple docstring"""







    super().__init__()
    if version.parse(torch.__version__) < version.parse("""1.9.0"""):
     _snake_case      :					Optional[int]					=   self._mish_python
    else:
     _snake_case      :					List[str]					=   nn.functional.mish

   def        UpperCamelCase_					(   self							:     str			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return input * torch.tanh(nn.functional.softplus(lowerCAmelCase))


   def        UpperCamelCase_					(   self							:     str			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return self.act(lowerCAmelCase)







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def        UpperCamelCase_					(   self							:     Optional[Any]			,						lowerCAmelCase							:     Tensor)			->			Tensor:

    """simple docstring"""







    return input







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def        UpperCamelCase_					(   self							:     List[Any]			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     List[Any]=0.707_107			,						lowerCAmelCase							:     Optional[Any]=0.282_095)			->			List[Any]:

    """simple docstring"""







    _snake_case      :					List[str]					=   (input - mu).div(sigma * math.sqrt(2.0))
    return 0.5 * (1.0 + torch.erf(lowerCAmelCase))







class 						snake_case  (		nn.Module					):



   '''simple docstring'''






   def        UpperCamelCase_					(   self							:     Any			,						lowerCAmelCase							:     Dict)			->			Union[str, Any]:

    """simple docstring"""







    _snake_case      :					Dict					=   nn.functional.relu(lowerCAmelCase)
    _snake_case      :					Tuple					=   torch.square(lowerCAmelCase)
    return squared







class 						snake_case  (		SCREAMING_SNAKE_CASE_					):



   '''simple docstring'''






   def __getitem__(   self							:     Dict			,						lowerCAmelCase							:     Tuple)			->			str:

    """simple docstring"""







    _snake_case      :					Optional[Any]					=   super().__getitem__(lowerCAmelCase)
    _snake_case		,      _snake_case      :					List[str]					=   content if isinstance(lowerCAmelCase			,						lowerCAmelCase) else (content, {})
    return cls(**lowerCAmelCase)
a__      =		{
    """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,
}
a__      =		ClassInstantier(ACTaCLS)



def  lowercase    (				SCREAMING_SNAKE_CASE__    :		Any			)					->       int:
 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
a__      =		get_activation("""gelu_python""")
a__      =		get_activation("""gelu_new""")
a__      =		get_activation("""gelu""")
a__      =		get_activation("""gelu_fast""")
a__      =		get_activation("""quick_gelu""")
a__      =		get_activation("""silu""")
a__      =		get_activation("""mish""")

a__      =		get_activation("""linear""")
