

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 		SCREAMING_SNAKE_CASE__				(							_UpperCAmelCase,      _UpperCAmelCase=0,      _UpperCAmelCase=None				)					-> Dict:



    '''simple docstring'''






    return [" ".join(x.split(' '				)[begin:end]				) for x in answers if len(x.split(' '				)				) > 1]
def 		SCREAMING_SNAKE_CASE__				(							_UpperCAmelCase,      _UpperCAmelCase=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 		SCREAMING_SNAKE_CASE__				(							_UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase=1,      _UpperCAmelCase=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 		SCREAMING_SNAKE_CASE__				(							_UpperCAmelCase				)					-> 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 		SCREAMING_SNAKE_CASE__				(							)					-> 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 _UpperCAmelCase				: (len(x['question']				) + len(x['context']				)) < 4 * 4_096				)
    lowerCAmelCase    			=      short_validation_dataset.filter(lambda _UpperCAmelCase				: 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(*_UpperCAmelCase,      **_UpperCAmelCase				):
        lowerCAmelCase    			=      model(*_UpperCAmelCase,      **_UpperCAmelCase				)
        return start_logits, end_logits, jnp.argmax(_UpperCAmelCase,      axis=-1				)

    def evaluate(_UpperCAmelCase				):
        # 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=4_096,      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 1_024
        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 _UpperCAmelCase				: x["match"] == 1				)				)
    print('EM score:',      (matched / total) * 100,      '%'				)


if __name__ == "__main__":
 main()








__A   :     List[Any]      =			'''
# Installation de Transformers
! pip install transformers datasets
# Pour installer à partir du code source au lieu de la dernière version, commentez la commande ci-dessus et décommentez la suivante.
# ! pip install git+https://github.com/huggingface/transformers.git
'''

__A   :     Dict      =			[{'''type''': '''code''', '''content''': INSTALL_CONTENT}]
__A   :     Tuple      =			{
    '''{processor_class}''': '''FakeProcessorClass''',
    '''{model_class}''': '''FakeModelClass''',
    '''{object_class}''': '''FakeObjectClass''',
}





