





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

from transformers import BigBirdTokenizerFast


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







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







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

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

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







def       lowerCAmelCase__(	lowercase       :    Any				) ->    Any:
  __snake_case   :				str        =       sample["question"]["text"]
  __snake_case   :				Tuple        =       sample["document"]["tokens"]["token"]
  __snake_case   :				Optional[Any]        =       sample["document"]["tokens"]["is_html"]
  __snake_case   :				Optional[int]        =       sample["annotations"]["long_answer"]
  __snake_case   :				str        =       sample["annotations"]["short_answers"]

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

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

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

  __snake_case   :				Any        =       "vasudevgupta/flax-bigbird-natural-questions"
  __snake_case   :				str        =       FlaxBigBirdForNaturalQuestions.from_pretrained(lowercase				)
  __snake_case   :				Optional[Any]        =       BigBirdTokenizerFast.from_pretrained(lowercase				)

  @jax.jit
  def forward(*lowercase       :    str   ,     **lowercase       :    Optional[int]				):
    __snake_case    ,   __snake_case    ,   __snake_case   :				int        =       model(*lowercase   ,     **lowercase				)
    return start_logits, end_logits, jnp.argmax(lowercase   ,     axis=-1				)

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

    __snake_case    ,   __snake_case    ,   __snake_case   :				int        =       forward(**lowercase				)

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

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

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

    __snake_case   :				Union[str, Any]        =       38 if predicted_category == "short" else 1024
    __snake_case    ,   __snake_case   :				Any        =       get_best_valid_start_end_idx(
        start_scores[0]   ,     end_scores[0]   ,     top_k=8   ,     max_size=lowercase				)

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

    __snake_case   :				Dict        =       expand_to_aliases(example["targets"]   ,     make_sub_answers=lowercase				)
    __snake_case   :				Optional[int]        =       expand_to_aliases(example["output"]				)

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

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

    return example

  __snake_case   :				Optional[int]        =       short_validation_dataset.map(lowercase				)

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


if __name__ == "__main__":
     main()








import torch
from torch import autograd




class 				_lowerCamelCase (    autograd.Function    ):



      """simple docstring"""






      @staticmethod
      def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   str:






        '''simple docstring'''






        __snake_case   :				Dict        =       inputs.numel()
        __snake_case   :				Optional[int]        =       int(0.005 * nb_elems   ) + 1
        if sigmoid:
          __snake_case   :				Any        =       (torch.sigmoid(UpperCAmelCase   ) > threshold).type(inputs.type()   )
        else:
          __snake_case   :				Optional[int]        =       (inputs > threshold).type(inputs.type()   )
        if mask.sum() < nb_min:
          # We limit the pruning so that at least 0.5% (half a percent) of the weights are remaining
          __snake_case   :				List[str]        =       inputs.flatten().kthvalue(max(nb_elems - nb_min				,						1   )   ).values
          __snake_case   :				Union[str, Any]        =       (inputs > k_threshold).type(inputs.type()   )
        return mask







      @staticmethod
      def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase   )       ->   Optional[int]:






        '''simple docstring'''






        return gradOutput, None, None




class 				_lowerCamelCase (    autograd.Function    ):



      """simple docstring"""






      @staticmethod
      def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   Optional[int]:






        '''simple docstring'''






        __snake_case   :				Optional[int]        =       inputs.clone()
        __snake_case    ,   __snake_case   :				int        =       inputs.flatten().sort(descending=UpperCAmelCase   )
        __snake_case   :				Tuple        =       int(threshold * inputs.numel()   )

        # flat_out and mask access the same memory.
        __snake_case   :				Optional[int]        =       mask.flatten()
        __snake_case   :				Union[str, Any]        =       0
        __snake_case   :				Any        =       1
        return mask







      @staticmethod
      def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase   )       ->   List[str]:






        '''simple docstring'''






        return gradOutput, None




class 				_lowerCamelCase (    a    ):



      """simple docstring"""






      @staticmethod
      def 						UpperCAmelCase		( UpperCAmelCase				,						UpperCAmelCase   )       ->   Any:






        '''simple docstring'''






        __snake_case   :				Union[str, Any]        =       inputs.clone()
        __snake_case    ,   __snake_case   :				Optional[Any]        =       inputs.abs().flatten().sort(descending=UpperCAmelCase   )
        __snake_case   :				List[Any]        =       int(threshold * inputs.numel()   )

        # flat_out and mask access the same memory.
        __snake_case   :				Optional[Any]        =       mask.flatten()
        __snake_case   :				Tuple        =       0
        __snake_case   :				Any        =       1
        return mask

