code: 
  
  
  
  
  
  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()
  
  
style_context: 
  
  
  
  
  
  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
  
  