




'''simple docstring'''





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

from transformers import BigBirdTokenizerFast


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





def 			__a			(     _UpperCamelCase:					Dict      ,		_UpperCamelCase:					Union[str, Any]=0      ,		_UpperCamelCase:					Optional[Any]=None					)		->				Any:

     """simple docstring"""




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





def 			__a			(     _UpperCamelCase:					Optional[int]      ,		_UpperCamelCase:					str=False					)		->				Tuple:

     """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"
          _snake_case							=   (
              given_answers + get_sub_answers(_UpperCamelCase      ,		begin=1					) + get_sub_answers(_UpperCamelCase      ,		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(_UpperCamelCase					)





def 			__a			(     _UpperCamelCase:					Tuple      ,		_UpperCamelCase:					Dict      ,		_UpperCamelCase:					int=1      ,		_UpperCamelCase:					Optional[int]=100					)		->				Any:

     """simple docstring"""




     _snake_case							,      _snake_case							=   jax.lax.top_k(_UpperCamelCase      ,		_UpperCamelCase					)
     _snake_case							,      _snake_case							=   jax.lax.top_k(_UpperCamelCase      ,		_UpperCamelCase					)

     _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(_UpperCamelCase					).item()

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





def 			__a			(     _UpperCamelCase:					Dict					)		->				int:

     """simple docstring"""




     _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(_UpperCamelCase					)					) 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(_UpperCamelCase					)					)

     _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(_UpperCamelCase					)					) if not html[i]]					)
          if new_answer not in long_targets:
               long_targets.append(_UpperCamelCase					)

     _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 			__a			(     )		->				Optional[Any]:

     """simple docstring"""




     _snake_case							=   load_from_disk("natural-questions-validation"					)
     _snake_case							=   dataset.map(_UpperCamelCase					).remove_columns(["annotations", "document", "id"]					)
     print(_UpperCamelCase					)

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

     _snake_case							=   "vasudevgupta/flax-bigbird-natural-questions"
     _snake_case							=   FlaxBigBirdForNaturalQuestions.from_pretrained(_UpperCamelCase					)
     _snake_case							=   BigBirdTokenizerFast.from_pretrained(_UpperCamelCase					)

     @jax.jit
     def forward(*_UpperCamelCase:					Any      ,		**_UpperCamelCase:					int					):
          _snake_case							,      _snake_case							,      _snake_case							=   model(*_UpperCamelCase      ,		**_UpperCamelCase					)
          return start_logits, end_logits, jnp.argmax(_UpperCamelCase      ,		axis=-1					)

     def evaluate(_UpperCamelCase:					Optional[Any]					):
          # 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=4_096      ,		padding="max_length"      ,		truncation=_UpperCamelCase      ,		)

          _snake_case							,      _snake_case							,      _snake_case							=   forward(**_UpperCamelCase					)

          _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 1_024
          _snake_case							,      _snake_case							=   get_best_valid_start_end_idx(
              start_scores[0]      ,		end_scores[0]      ,		top_k=8      ,		max_size=_UpperCamelCase					)

          _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=_UpperCamelCase					)
          _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(_UpperCamelCase					) > 0

          return example

     _snake_case							=   short_validation_dataset.map(_UpperCamelCase					)

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


if __name__ == "__main__":
   main()








'''simple docstring'''





import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union

class      _a       (		enum.Enum							):
     SCREAMING_SNAKE_CASE_							:					Optional[int]					       =		"""timm"""
     SCREAMING_SNAKE_CASE_							:					int					       =		"""transformers"""







def 			__a			(     _UpperCamelCase:					Optional[Iterable[str]]      ,		_UpperCamelCase:					Optional[Iterable[int]]      ,		_UpperCamelCase:					Optional[Iterable[str]]					)		->				Union[str, Any]:

     """simple docstring"""




     if stage_names is None:
          raise ValueError("Stage_names must be set for transformers backbones"					)

     if out_features is not None:
          if not isinstance(_UpperCamelCase      ,		(list,)					):
               raise ValueError(F"""out_features must be a list {type(_UpperCamelCase					)}"""					)
          if any(feat not in stage_names for feat in out_features					):
               raise ValueError(F"""out_features must be a subset of stage_names: {stage_names} got {out_features}"""					)

     if out_indices is not None:
          if not isinstance(_UpperCamelCase      ,		(list, tuple)					):
               raise ValueError(F"""out_indices must be a list or tuple, got {type(_UpperCamelCase					)}"""					)
          if any(idx >= len(_UpperCamelCase					) for idx in out_indices					):
               raise ValueError("out_indices must be valid indices for stage_names {stage_names}, got {out_indices}"					)

     if out_features is not None and out_indices is not None:
          if len(_UpperCamelCase					) != len(_UpperCamelCase					):
               raise ValueError("out_features and out_indices should have the same length if both are set"					)
          if out_features != [stage_names[idx] for idx in out_indices]:
               raise ValueError("out_features and out_indices should correspond to the same stages if both are set"					)





def 			__a			(     _UpperCamelCase:					Optional[List[str]]      ,		_UpperCamelCase:					Optional[Union[List[int], Tuple[int]]]      ,		_UpperCamelCase:					List[str]      ,		)		->				Any:

     """simple docstring"""




     if out_indices is None and out_features is None:
          _snake_case							=   [len(_UpperCamelCase					) - 1]
          _snake_case							=   [stage_names[-1]]
     elif out_indices is None and out_features is not None:
          _snake_case							=   [stage_names.index(_UpperCamelCase					) for layer in out_features]
     elif out_features is None and out_indices is not None:
          _snake_case							=   [stage_names[idx] for idx in out_indices]
     return out_features, out_indices





def 			__a			(     _UpperCamelCase:					Optional[List[str]]      ,		_UpperCamelCase:					Optional[Union[List[int], Tuple[int]]]      ,		_UpperCamelCase:					List[str]      ,		)		->				Tuple[List[str], List[int]]:

     """simple docstring"""




     verify_out_features_out_indices(out_features=_UpperCamelCase      ,		out_indices=_UpperCamelCase      ,		stage_names=_UpperCamelCase					)
     _snake_case							,      _snake_case							=   _align_output_features_output_indices(
         out_features=_UpperCamelCase      ,		out_indices=_UpperCamelCase      ,		stage_names=_UpperCamelCase					)
     # Verify that the aligned out_features and out_indices are valid
     verify_out_features_out_indices(out_features=_UpperCamelCase      ,		out_indices=_UpperCamelCase      ,		stage_names=_UpperCamelCase					)
     return output_features, output_indices

class      _a       :
     SCREAMING_SNAKE_CASE_							:					Optional[BackboneType]					       =		None






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    None:
          if getattr(self      ,"_backbone"      ,_SCREAMING_SNAKE_CASE					) is None:
               raise ValueError("self._backbone must be set before calling _init_timm_backbone"					)

          # These will diagree with the defaults for the transformers models e.g. for resnet50
          # the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4']
          # the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4']
          _snake_case							=   [stage["module"] for stage in self._backbone.feature_info.info]
          _snake_case							=   [stage["num_chs"] for stage in self._backbone.feature_info.info]
          _snake_case							=   self._backbone.feature_info.out_indices
          _snake_case							=   self._backbone.feature_info.module_name()

          # We verify the out indices and out features are valid
          verify_out_features_out_indices(
              out_features=_SCREAMING_SNAKE_CASE      ,out_indices=_SCREAMING_SNAKE_CASE      ,stage_names=self.stage_names					)
          _snake_case							,      _snake_case							=   out_features, out_indices






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    None:
          _snake_case							=   getattr(_SCREAMING_SNAKE_CASE      ,"stage_names"					)
          _snake_case							=   getattr(_SCREAMING_SNAKE_CASE      ,"out_features"      ,_SCREAMING_SNAKE_CASE					)
          _snake_case							=   getattr(_SCREAMING_SNAKE_CASE      ,"out_indices"      ,_SCREAMING_SNAKE_CASE					)

          _snake_case							=   stage_names
          _snake_case							,      _snake_case							=   get_aligned_output_features_output_indices(
              out_features=_SCREAMING_SNAKE_CASE      ,out_indices=_SCREAMING_SNAKE_CASE      ,stage_names=_SCREAMING_SNAKE_CASE					)
          # Number of channels for each stage. This is set in the transformer backbone model init
          _snake_case							=   None






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    None:
          _snake_case							=   config

          _snake_case							=   getattr(_SCREAMING_SNAKE_CASE      ,"use_timm_backbone"      ,_SCREAMING_SNAKE_CASE					)
          _snake_case							=   BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

          if self.backbone_type == BackboneType.TIMM:
               self._init_timm_backbone(_SCREAMING_SNAKE_CASE					)
          elif self.backbone_type == BackboneType.TRANSFORMERS:
               self._init_transformers_backbone(_SCREAMING_SNAKE_CASE					)
          else:
               raise ValueError(f"""backbone_type {self.backbone_type} not supported."""					)






     @property
     def        _lowercase			(				self					)							->    int:
          return self._out_features






     @out_features.setter
     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    str:
          _snake_case							,      _snake_case							=   get_aligned_output_features_output_indices(
              out_features=_SCREAMING_SNAKE_CASE      ,out_indices=_SCREAMING_SNAKE_CASE      ,stage_names=self.stage_names					)






     @property
     def        _lowercase			(				self					)							->    Any:
          return self._out_indices






     @out_indices.setter
     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    int:
          _snake_case							,      _snake_case							=   get_aligned_output_features_output_indices(
              out_features=_SCREAMING_SNAKE_CASE      ,out_indices=_SCREAMING_SNAKE_CASE      ,stage_names=self.stage_names					)






     @property
     def        _lowercase			(				self					)							->    List[str]:
          # the current backbones will output the number of channels for each stage
          # even if that stage is not in the out_features list.
          return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names					)}






     @property
     def        _lowercase			(				self					)							->    Any:
          return [self.out_feature_channels[name] for name in self.out_features]






     def        _lowercase			(				self      ,*_SCREAMING_SNAKE_CASE      ,**_SCREAMING_SNAKE_CASE					)							->    str:
          _snake_case							=   dict(inspect.signature(self.forward					).parameters					)
          _snake_case							=   {k: v for k, v in kwargs.items() if k in signature}
          return self(*_SCREAMING_SNAKE_CASE      ,**_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE = None      ,_SCREAMING_SNAKE_CASE = None      ,_SCREAMING_SNAKE_CASE = None      ,)							->    Any:
          raise NotImplementedError("This method should be implemented by the derived class."					)






     def        _lowercase			(				self					)							->    Tuple:
          _snake_case							=   super().to_dict()
          _snake_case							=   output.pop("_out_features"					)
          _snake_case							=   output.pop("_out_indices"					)
          return output







class      _a       :






     @property
     def        _lowercase			(				self					)							->    List[str]:
          return self._out_features






     @out_features.setter
     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    str:
          _snake_case							,      _snake_case							=   get_aligned_output_features_output_indices(
              out_features=_SCREAMING_SNAKE_CASE      ,out_indices=_SCREAMING_SNAKE_CASE      ,stage_names=self.stage_names					)






     @property
     def        _lowercase			(				self					)							->    Any:
          return self._out_indices






     @out_indices.setter
     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    Optional[Any]:
          _snake_case							,      _snake_case							=   get_aligned_output_features_output_indices(
              out_features=_SCREAMING_SNAKE_CASE      ,out_indices=_SCREAMING_SNAKE_CASE      ,stage_names=self.stage_names					)






     def        _lowercase			(				self					)							->    List[str]:
          _snake_case							=   super().to_dict()
          _snake_case							=   output.pop("_out_features"					)
          _snake_case							=   output.pop("_out_indices"					)
          return output


