



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    (					__lowerCamelCase :  Union[str, Any]		,						__lowerCamelCase :  List[str]=0		,						__lowerCamelCase :  Optional[int]=None  ):



 """simple docstring"""





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






def 						_lowerCAmelCase    (					__lowerCamelCase :  Optional[int]		,						__lowerCamelCase :  List[Any]=False  ):



 """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"
  __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						(
      given_answers + get_sub_answers(__lowerCamelCase		,						begin=1  ) + get_sub_answers(__lowerCamelCase		,						end=-1  )
  )
 __SCREAMING_SNAKE_CASE							:     Tuple							=						[]
 for answer in given_answers:
  __SCREAMING_SNAKE_CASE							:     Any							=						answer.replace("_"		,						" "  ).lower()
  __SCREAMING_SNAKE_CASE							:     str							=						"".join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else " " for c in alias  )
  answers.append(" ".join(alias.split()  ).strip()  )
 return set(__lowerCamelCase  )






def 						_lowerCAmelCase    (					__lowerCamelCase :  Optional[int]		,						__lowerCamelCase :  Dict		,						__lowerCamelCase :  List[str]=1		,						__lowerCamelCase :  int=100  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     Any							=						jax.lax.top_k(__lowerCamelCase		,						__lowerCamelCase  )
 __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						jax.lax.top_k(__lowerCamelCase		,						__lowerCamelCase  )

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

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






def 						_lowerCAmelCase    (					__lowerCamelCase :  Optional[int]  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[str]							=						sample["question"]["text"]
 __SCREAMING_SNAKE_CASE							:     str							=						sample["document"]["tokens"]["token"]
 __SCREAMING_SNAKE_CASE							:     Dict							=						sample["document"]["tokens"]["is_html"]
 __SCREAMING_SNAKE_CASE							:     str							=						sample["annotations"]["long_answer"]
 __SCREAMING_SNAKE_CASE							:     List[Any]							=						sample["annotations"]["short_answers"]

 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						" ".join([context[i] for i in range(len(__lowerCamelCase  )  ) 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",
   }

 __SCREAMING_SNAKE_CASE							:     Dict							=						[]
 for s in short_answers:
  short_targets.extend(s["text"]  )
 __SCREAMING_SNAKE_CASE							:     Tuple							=						list(set(__lowerCamelCase  )  )

 __SCREAMING_SNAKE_CASE							:     str							=						[]
 for s in long_answers:
  if s["start_token"] == -1:
   continue
  __SCREAMING_SNAKE_CASE							:     Any							=						context[s["start_token"] : s["end_token"]]
  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						is_html[s["start_token"] : s["end_token"]]
  __SCREAMING_SNAKE_CASE							:     List[str]							=						" ".join([answer[i] for i in range(len(__lowerCamelCase  )  ) if not html[i]]  )
  if new_answer not in long_targets:
   long_targets.append(__lowerCamelCase  )

 __SCREAMING_SNAKE_CASE							:     Dict							=						"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    (					):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[Any]							=						load_from_disk("natural-questions-validation"  )
 __SCREAMING_SNAKE_CASE							:     List[str]							=						dataset.map(__lowerCamelCase  ).remove_columns(["annotations", "document", "id"]  )
 print(__lowerCamelCase  )

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

 __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						"vasudevgupta/flax-bigbird-natural-questions"
 __SCREAMING_SNAKE_CASE							:     Tuple							=						FlaxBigBirdForNaturalQuestions.from_pretrained(__lowerCamelCase  )
 __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						BigBirdTokenizerFast.from_pretrained(__lowerCamelCase  )

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

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

  __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     Optional[Any]							=						forward(**__lowerCamelCase  )

  __SCREAMING_SNAKE_CASE							:     Tuple							=						CATEGORY_MAPPING[category.item()]

  __SCREAMING_SNAKE_CASE							:     Any							=						example["long"] + example["short"]
  if example["category"] in ["yes", "no", "null"]:
   __SCREAMING_SNAKE_CASE							:     List[str]							=						[example["category"]]
  __SCREAMING_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"]:
   __SCREAMING_SNAKE_CASE							:     Any							=						[predicted_category]
   __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						example["output"] == example["targets"]
   __SCREAMING_SNAKE_CASE							:     List[str]							=						predicted_category != "null"
   return example

  __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						38 if predicted_category == "short" else 1024
  __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     str							=						get_best_valid_start_end_idx(
      start_scores[0]		,						end_scores[0]		,						top_k=8		,						max_size=__lowerCamelCase  )

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

  __SCREAMING_SNAKE_CASE							:     Dict							=						expand_to_aliases(example["targets"]		,						make_sub_answers=__lowerCamelCase  )
  __SCREAMING_SNAKE_CASE							:     str							=						expand_to_aliases(example["output"]  )

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

  # if there is a common element, it's a exact match
  __SCREAMING_SNAKE_CASE							:     Any							=						len(list(answers & predictions  )  ) > 0
  __SCREAMING_SNAKE_CASE							:     Optional[int]							=						predicted_category != "null" and len(__lowerCamelCase  ) > 0

  return example

 __SCREAMING_SNAKE_CASE							:     Any							=						short_validation_dataset.map(__lowerCamelCase  )

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


if __name__ == "__main__":
      main()











import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


if is_tpu_available(check_device=False):
      import torch_xla.core.xla_model as xm






def 						_lowerCAmelCase    (					__lowerCamelCase :  Any		,						__lowerCamelCase :  str  ):



 """simple docstring"""





 if isinstance(__lowerCamelCase		,						(list, tuple)  ):
  return honor_type(__lowerCamelCase		,						(move_to_device(__lowerCamelCase		,						__lowerCamelCase  ) for t in state)  )
 elif isinstance(__lowerCamelCase		,						__lowerCamelCase  ):
  return type(__lowerCamelCase  )({k: move_to_device(__lowerCamelCase		,						__lowerCamelCase  ) for k, v in state.items()}  )
 elif isinstance(__lowerCamelCase		,						torch.Tensor  ):
  return state.to(__lowerCamelCase  )
 return state






class    _SCREAMING_SNAKE_CASE						(torch.optim.Optimizer			):





     def __init__(							self			:				int  ,	UpperCamelCase			:				Any  ,	UpperCamelCase			:				List[str]=True  ,	UpperCamelCase			:				List[Any]=None					)->Optional[int]:
      __SCREAMING_SNAKE_CASE							:     str							=						optimizer
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						scaler
      __SCREAMING_SNAKE_CASE							:     Any							=						AcceleratorState()
      __SCREAMING_SNAKE_CASE							:     int							=						GradientState()
      __SCREAMING_SNAKE_CASE							:     str							=						device_placement
      __SCREAMING_SNAKE_CASE							:     int							=						False
      __SCREAMING_SNAKE_CASE							:     Dict							=						None

      # Handle device placement
      if device_placement:
       __SCREAMING_SNAKE_CASE							:     List[str]							=						self.optimizer.state_dict()
       if self.accelerator_state.distributed_type == DistributedType.TPU:
        xm.send_cpu_data_to_device(UpperCamelCase  ,	self.accelerator_state.device					)
       else:
        __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						move_to_device(UpperCamelCase  ,	self.accelerator_state.device					)
       self.optimizer.load_state_dict(UpperCamelCase					)





     @property
     def 						__snake_case     (							self			:				Optional[Any]					)->List[str]:
      return self.optimizer.state





     @state.setter
     def 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				List[str]					)->Dict:
      __SCREAMING_SNAKE_CASE							:     List[str]							=						state





     @property
     def 						__snake_case     (							self			:				List[Any]					)->str:
      return self.optimizer.param_groups





     @param_groups.setter
     def 						__snake_case     (							self			:				List[str]  ,	UpperCamelCase			:				Any					)->Dict:
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						param_groups





     @property
     def 						__snake_case     (							self			:				str					)->str:
      return self.optimizer.defaults





     @defaults.setter
     def 						__snake_case     (							self			:				int  ,	UpperCamelCase			:				Tuple					)->List[str]:
      __SCREAMING_SNAKE_CASE							:     List[str]							=						defaults





     def 						__snake_case     (							self			:				Any  ,	UpperCamelCase			:				Tuple					)->Optional[int]:
      self.optimizer.add_param_group(UpperCamelCase					)





     def 						__snake_case     (							self			:				Union[str, Any]  ,	UpperCamelCase			:				Union[str, Any]					)->Dict:
      if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
       xm.send_cpu_data_to_device(UpperCamelCase  ,	self.accelerator_state.device					)
      self.optimizer.load_state_dict(UpperCamelCase					)





     def 						__snake_case     (							self			:				List[str]					)->List[Any]:
      return self.optimizer.state_dict()





     def 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				Any=None					)->Any:
      if self.gradient_state.sync_gradients:
       __SCREAMING_SNAKE_CASE							:     Optional[int]							=						"set_to_none" in inspect.signature(self.optimizer.zero_grad					).parameters
       if accept_arg:
        if set_to_none is None:
         __SCREAMING_SNAKE_CASE							:     Optional[int]							=						False
        self.optimizer.zero_grad(set_to_none=UpperCamelCase					)
       else:
        if set_to_none is not None:
         raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer."					)
        self.optimizer.zero_grad()





     def 						__snake_case     (							self			:				Dict  ,	UpperCamelCase			:				Tuple=None					)->List[Any]:
      if self.gradient_state.sync_gradients:
       if self.accelerator_state.distributed_type == DistributedType.TPU:
        __SCREAMING_SNAKE_CASE							:     List[str]							=						{"closure": closure} if closure is not None else {}
        xm.optimizer_step(self.optimizer  ,	optimizer_args=UpperCamelCase					)
       elif self.scaler is not None:
        __SCREAMING_SNAKE_CASE							:     Optional[int]							=						False
        if self._last_scale is None:
         # `get_scale` is an async operation requiring full synchronization
         # on CPU and GPUs before finishing. As a result, we store away
         # the prior one to reduce the call overhead
         __SCREAMING_SNAKE_CASE							:     List[Any]							=						self.scaler.get_scale()
         __SCREAMING_SNAKE_CASE							:     List[Any]							=						True
        self.scaler.step(self.optimizer  ,	UpperCamelCase					)
        self.scaler.update()
        __SCREAMING_SNAKE_CASE							:     Dict							=						self.scaler.get_scale()
        if not new_scale:
         # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
         __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						scale_after < self._last_scale
        __SCREAMING_SNAKE_CASE							:     List[Any]							=						scale_after

       else:
        self.optimizer.step(UpperCamelCase					)





     def 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				Optional[Any]					)->Any:
      for param_group in self.optimizer.param_groups:
       __SCREAMING_SNAKE_CASE							:     Tuple							=						[parameters_map.get(UpperCamelCase  ,	UpperCamelCase					) for p in param_group["params"]]





     @property
     def 						__snake_case     (							self			:				List[Any]					)->int:
      warnings.warn(
          "The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use "
          "`optimizer.step_was_skipped` instead."  ,	UpperCamelCase  ,	)
      return self._is_overflow





     @property
     def 						__snake_case     (							self			:				Optional[int]					)->List[Any]:
      return self._is_overflow





     def __getstate__(							self			:				List[str]					)->List[str]:
      return self.__dict__.copy()






     def __setstate__(							self			:				Union[str, Any]  ,	UpperCamelCase			:				str					)->int:
      self.__dict__.update(UpperCamelCase					)






