code: 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(				__snake_case						,						__snake_case=0						,						__snake_case=None	)  ->   Tuple:
    return [" ".join(x.split(" "	)[begin:end]	) for x in answers if len(x.split(" "	)	) > 1]
  
  
  
  def 							_lowerCamelCase(				__snake_case						,						__snake_case=False	)  ->   List[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		       =     (
          given_answers + get_sub_answers(__A						,						begin=1	) + get_sub_answers(__A						,						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(__A	)
  
  
  
  def 							_lowerCamelCase(				__snake_case						,						__snake_case						,						__snake_case=1						,						__snake_case=100	)  ->   str:
    __snake_case       ,       __snake_case		       =     jax.lax.top_k(__A						,						__A	)
    __snake_case       ,       __snake_case		       =     jax.lax.top_k(__A						,						__A	)
  
    __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(__A	).item()
  
    return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]
  
  
  
  def 							_lowerCamelCase(				__snake_case	)  ->   str:
    __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(__A	)	) 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(__A	)	)
  
    __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(__A	)	) if not html[i]]	)
      if new_answer not in long_targets:
        long_targets.append(__A	)
  
    __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 							_lowerCamelCase(				)  ->   Optional[Any]:
    __snake_case		       =     load_from_disk("natural-questions-validation"	)
    __snake_case		       =     dataset.map(__A	).remove_columns(["annotations", "document", "id"]	)
    print(__A	)
  
    __snake_case		       =     dataset.filter(lambda __snake_case	: (len(x["question"]	) + len(x["context"]	)) < 4 * 4096	)
    __snake_case		       =     short_validation_dataset.filter(lambda __snake_case	: x["category"] != "null"	)
    short_validation_dataset
  
    __snake_case		       =     "vasudevgupta/flax-bigbird-natural-questions"
    __snake_case		       =     FlaxBigBirdForNaturalQuestions.from_pretrained(__A	)
    __snake_case		       =     BigBirdTokenizerFast.from_pretrained(__A	)
  
    @jax.jit
    def forward(*__snake_case						,						**__snake_case	):
      __snake_case       ,       __snake_case       ,       __snake_case		       =     model(*__A						,						**__A	)
      return start_logits, end_logits, jnp.argmax(__A						,						axis=-1	)
  
    def evaluate(__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=__A						,						)
  
      __snake_case       ,       __snake_case       ,       __snake_case		       =     forward(**__A	)
  
      __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=__A	)
  
      __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=__A	)
      __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(__A	) > 0
  
      return example
  
    __snake_case		       =     short_validation_dataset.map(__A	)
  
    __snake_case		       =     len(__A	)
    __snake_case		       =     len(short_validation_dataset.filter(lambda __snake_case	: x["match"] == 1	)	)
    print("EM score:"						,						(matched / total) * 100						,						"%"	)
  
  
  if __name__ == "__main__":
    main()
  
  
  
  
  
  
style_context: import posixpath
  from pathlib import Path
  from unittest.mock import patch
  
  import pytest
  from fsspec.implementations.local import AbstractFileSystem, LocalFileSystem, stringify_path
  from fsspec.registry import _registry as _fsspec_registry
  
  
  
  
  class 					UpperCamelCase	(     snake_case__      ):
  					__UpperCamelCase						=   """mock"""
  
  
  
  
  					def __init__( self :					List[Any]      ,*_lowerCAmelCase :					Optional[int]      ,_lowerCAmelCase :					Union[str, Any]      ,**_lowerCAmelCase :					Union[str, Any]						):
  
  
  							"""simple docstring"""
  
  
  
  
  							super().__init__()
  							__snake_case		       =     LocalFileSystem(*_lowerCAmelCase      ,**_lowerCAmelCase						)
  							__snake_case		       =     Path(_lowerCAmelCase						).resolve().as_posix() + "/"
  
  
  
  
  					def 							UpperCamelCase_						( self :					List[Any]      ,_lowerCAmelCase :					Optional[Any]      ,*_lowerCAmelCase :					Union[str, Any]      ,**_lowerCAmelCase :					Dict						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs.mkdir(_lowerCAmelCase      ,*_lowerCAmelCase      ,**_lowerCAmelCase						)
  
  
  
  
  					def 							UpperCamelCase_						( self :					int      ,_lowerCAmelCase :					Tuple      ,*_lowerCAmelCase :					Optional[Any]      ,**_lowerCAmelCase :					Optional[Any]						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs.makedirs(_lowerCAmelCase      ,*_lowerCAmelCase      ,**_lowerCAmelCase						)
  
  
  
  
  					def 							UpperCamelCase_						( self :					Tuple      ,_lowerCAmelCase :					Any						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs.rmdir(_lowerCAmelCase						)
  
  
  
  
  					def 							UpperCamelCase_						( self :					List[Any]      ,_lowerCAmelCase :					List[Any]      ,_lowerCAmelCase :					str=True      ,*_lowerCAmelCase :					int      ,**_lowerCAmelCase :					List[str]						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							__snake_case		       =     self._fs.ls(_lowerCAmelCase      ,detail=_lowerCAmelCase      ,*_lowerCAmelCase      ,**_lowerCAmelCase						)
  							if detail:
  									return [{**info, "name": info["name"][len(self.local_root_dir						) :]} for info in out]
  							else:
  									return [name[len(self.local_root_dir						) :] for name in out]
  
  
  
  
  					def 							UpperCamelCase_						( self :					List[str]      ,_lowerCAmelCase :					Any      ,*_lowerCAmelCase :					Any      ,**_lowerCAmelCase :					str						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							__snake_case		       =     dict(self._fs.info(_lowerCAmelCase      ,*_lowerCAmelCase      ,**_lowerCAmelCase						)						)
  							__snake_case		       =     out["name"][len(self.local_root_dir						) :]
  							return out
  
  
  
  
  					def 							UpperCamelCase_						( self :					Tuple      ,_lowerCAmelCase :					Optional[int]      ,_lowerCAmelCase :					Optional[Any]      ,*_lowerCAmelCase :					int      ,**_lowerCAmelCase :					Optional[int]						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs.cp_file(_lowerCAmelCase      ,_lowerCAmelCase      ,*_lowerCAmelCase      ,**_lowerCAmelCase						)
  
  
  
  
  					def 							UpperCamelCase_						( self :					Union[str, Any]      ,_lowerCAmelCase :					int      ,*_lowerCAmelCase :					List[Any]      ,**_lowerCAmelCase :					int						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs.rm_file(_lowerCAmelCase      ,*_lowerCAmelCase      ,**_lowerCAmelCase						)
  
  
  
  
  					def 							UpperCamelCase_						( self :					Optional[Any]      ,_lowerCAmelCase :					str      ,*_lowerCAmelCase :					Optional[int]      ,**_lowerCAmelCase :					int						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs.rm(_lowerCAmelCase      ,*_lowerCAmelCase      ,**_lowerCAmelCase						)
  
  
  
  
  					def 							UpperCamelCase_						( self :					int      ,_lowerCAmelCase :					Optional[int]      ,*_lowerCAmelCase :					str      ,**_lowerCAmelCase :					Optional[Any]						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs._open(_lowerCAmelCase      ,*_lowerCAmelCase      ,**_lowerCAmelCase						)
  
  
  
  
  					def 							UpperCamelCase_						( self :					Union[str, Any]      ,_lowerCAmelCase :					Tuple						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs.created(_lowerCAmelCase						)
  
  
  
  
  					def 							UpperCamelCase_						( self :					Union[str, Any]      ,_lowerCAmelCase :					int						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     posixpath.join(self.local_root_dir      ,self._strip_protocol(_lowerCAmelCase						)						)
  							return self._fs.modified(_lowerCAmelCase						)
  
  
  
  
  
  
  
  					@classmethod
  					def 							UpperCamelCase_						( cls :					Dict      ,_lowerCAmelCase :					Dict						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     stringify_path(_lowerCAmelCase						)
  							if path.startswith("mock://"						):
  									__snake_case		       =     path[7:]
  							return path
  
  
  
  
  class 					UpperCamelCase	(     snake_case__      ):
  					__UpperCamelCase						=   """tmp"""
  					__UpperCamelCase						=   None
  
  
  
  
  					def __init__( self :					int      ,*_lowerCAmelCase :					Tuple      ,**_lowerCAmelCase :					Union[str, Any]						):
  
  
  							"""simple docstring"""
  
  
  
  
  							assert self.tmp_dir is not None, "TmpDirFileSystem.tmp_dir is not set"
  							super().__init__(*_lowerCAmelCase      ,**_lowerCAmelCase      ,local_root_dir=self.tmp_dir      ,auto_mkdir=_lowerCAmelCase						)
  
  
  
  
  
  
  
  					@classmethod
  					def 							UpperCamelCase_						( cls :					int      ,_lowerCAmelCase :					str						):
  
  
  							"""simple docstring"""
  
  
  
  
  							__snake_case		       =     stringify_path(_lowerCAmelCase						)
  							if path.startswith("tmp://"						):
  									__snake_case		       =     path[6:]
  							return path
  
  
  
  
  @pytest.fixture
  def 							_lowerCamelCase(				)  ->   str:
  		__snake_case		       =     MockFileSystem
  		__snake_case		       =     TmpDirFileSystem
  		yield
  		del _fsspec_registry["mock"]
  		del _fsspec_registry["tmp"]
  
  
  
  @pytest.fixture
  def 							_lowerCamelCase(				__snake_case						,						__snake_case	)  ->   Dict:
  		__snake_case		       =     tmp_path_factory.mktemp("mockfs"	)
  		return MockFileSystem(local_root_dir=__snake_case						,						auto_mkdir=__snake_case	)
  
  
  
  
  
  @pytest.fixture
  def 							_lowerCamelCase(				__snake_case						,						__snake_case	)  ->   List[Any]:
  		__snake_case		       =     tmp_path_factory.mktemp("tmpfs"	)
  		with patch.object(__snake_case						,						"tmp_dir"						,						__snake_case	):
  				yield TmpDirFileSystem()
  
  
  
  
  
  