




'''simple docstring'''




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

from transformers import BigBirdTokenizerFast


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





def   a_							(   __snake_case							:    List[str]     , __snake_case							:    int=0     , __snake_case							:    int=None     )							->		Dict:



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





def   a_							(   __snake_case							:    Any     , __snake_case							:    Optional[Any]=False     )							->		int:



       """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"
              lowerCamelCase_							=(
                  given_answers + get_sub_answers(_UpperCAmelCase     , begin=1     ) + get_sub_answers(_UpperCAmelCase     , end=-1     )
              )
       lowerCamelCase_							=[]
       for answer in given_answers:
              lowerCamelCase_							=answer.replace('''_'''     , ''' '''     ).lower()
              lowerCamelCase_							=''.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_							(   __snake_case							:    List[Any]     , __snake_case							:    Any     , __snake_case							:    Any=1     , __snake_case							:    Any=100     )							->		Tuple:



       """simple docstring"""
       lowerCamelCase_							=jax.lax.top_k(_UpperCAmelCase     , _UpperCAmelCase     )
       lowerCamelCase_							=jax.lax.top_k(_UpperCAmelCase     , _UpperCAmelCase     )

       lowerCamelCase_							=best_end_idx[:, None] - best_start_idx[None, :]
       lowerCamelCase_							=jnp.logical_or(widths < 0     , widths > max_size     )
       lowerCamelCase_							=(best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
       lowerCamelCase_							=jnp.argmax(_UpperCAmelCase     ).item()

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





def   a_							(   __snake_case							:    Optional[Any]     )							->		Optional[int]:



       """simple docstring"""
       lowerCamelCase_							=sample['question']['text']
       lowerCamelCase_							=sample['document']['tokens']['token']
       lowerCamelCase_							=sample['document']['tokens']['is_html']
       lowerCamelCase_							=sample['annotations']['long_answer']
       lowerCamelCase_							=sample['annotations']['short_answers']

       lowerCamelCase_							=' '.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",
                     }

       lowerCamelCase_							=[]
       for s in short_answers:
              short_targets.extend(s['''text''']     )
       lowerCamelCase_							=list(set(_UpperCAmelCase     )     )

       lowerCamelCase_							=[]
       for s in long_answers:
              if s["start_token"] == -1:
                     continue
              lowerCamelCase_							=context[s['start_token'] : s['end_token']]
              lowerCamelCase_							=is_html[s['start_token'] : s['end_token']]
              lowerCamelCase_							=' '.join([answer[i] for i in range(len(_UpperCAmelCase     )     ) if not html[i]]     )
              if new_answer not in long_targets:
                     long_targets.append(_UpperCAmelCase     )

       lowerCamelCase_							='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_							(   )							->		Union[str, Any]:



       """simple docstring"""
       lowerCamelCase_							=load_from_disk('''natural-questions-validation'''     )
       lowerCamelCase_							=dataset.map(_UpperCAmelCase     ).remove_columns(['''annotations''', '''document''', '''id''']     )
       print(_UpperCAmelCase     )

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

       lowerCamelCase_							='vasudevgupta/flax-bigbird-natural-questions'
       lowerCamelCase_							=FlaxBigBirdForNaturalQuestions.from_pretrained(_UpperCAmelCase     )
       lowerCamelCase_							=BigBirdTokenizerFast.from_pretrained(_UpperCAmelCase     )

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

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

              lowerCamelCase_							=forward(**_UpperCAmelCase     )

              lowerCamelCase_							=CATEGORY_MAPPING[category.item()]

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

              if predicted_category in ["yes", "no", "null"]:
                     lowerCamelCase_							=[predicted_category]
                     lowerCamelCase_							=example['output'] == example['targets']
                     lowerCamelCase_							=predicted_category != 'null'
                     return example

              lowerCamelCase_							=38 if predicted_category == 'short' else 1024
              lowerCamelCase_							=get_best_valid_start_end_idx(
                  start_scores[0]     , end_scores[0]     , top_k=8     , max_size=_UpperCAmelCase     )

              lowerCamelCase_							=inputs['input_ids'][0].tolist()
              lowerCamelCase_							=[tokenizer.decode(input_ids[start_score : end_score + 1]     )]

              lowerCamelCase_							=expand_to_aliases(example['''targets''']     , make_sub_answers=_UpperCAmelCase     )
              lowerCamelCase_							=expand_to_aliases(example['''output''']     )

              # some preprocessing to both prediction and answer
              lowerCamelCase_							={''.join(a.split()     ) for a in answers}
              lowerCamelCase_							={''.join(p.split()     ) for p in predictions}
              lowerCamelCase_							={s for s in predictions if s not in ['``', '\'\'', '`', '\'']}

              # if there is a common element, it's a exact match
              lowerCamelCase_							=len(list(answers & predictions     )     ) > 0
              lowerCamelCase_							=predicted_category != 'null' and len(_UpperCAmelCase     ) > 0

              return example

       lowerCamelCase_							=short_validation_dataset.map(_UpperCAmelCase     )

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


if __name__ == "__main__":
     main()










'''simple docstring'''




import warnings

from .generation import FlaxGenerationMixin


class 						__UpperCamelCase (				lowerCamelCase__       ):
       # warning at import time
       warnings.warn(
           'Importing `FlaxGenerationMixin` from `src/transformers/generation_flax_utils.py` is deprecated and will '
           'be removed in Transformers v5. Import as `from transformers import FlaxGenerationMixin` instead.'    ,      lowerCamelCase__    ,      )




