




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

from transformers import BigBirdTokenizerFast


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




def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_=0    ,   lowercase_=None					)				->		Any:

    """simple docstring"""


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




def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_=False					)				->		Optional[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"
        __UpperCamelCase     						=  (
            given_answers + get_sub_answers(__UpperCamelCase    ,   begin=1					) + get_sub_answers(__UpperCamelCase    ,   end=-1					)
        )
    __UpperCamelCase     						=  []
    for answer in given_answers:
        __UpperCamelCase     						=  answer.replace('''_'''    ,   ''' '''					).lower()
        __UpperCamelCase     						=  ''''''.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      (     lowercase_    ,   lowercase_    ,   lowercase_=1    ,   lowercase_=1_00					)				->		Optional[int]:

    """simple docstring"""


    __UpperCamelCase						,					__UpperCamelCase     						=  jax.lax.top_k(__UpperCamelCase    ,   __UpperCamelCase					)
    __UpperCamelCase						,					__UpperCamelCase     						=  jax.lax.top_k(__UpperCamelCase    ,   __UpperCamelCase					)

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

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




def     __SCREAMING_SNAKE_CASE      (     lowercase_					)				->		Tuple:

    """simple docstring"""


    __UpperCamelCase     						=  sample['''question''']['''text''']
    __UpperCamelCase     						=  sample['''document''']['''tokens''']['''token''']
    __UpperCamelCase     						=  sample['''document''']['''tokens''']['''is_html''']
    __UpperCamelCase     						=  sample['''annotations''']['''long_answer''']
    __UpperCamelCase     						=  sample['''annotations''']['''short_answers''']

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

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

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

    __UpperCamelCase     						=  '''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"""


    __UpperCamelCase     						=  load_from_disk('''natural-questions-validation'''					)
    __UpperCamelCase     						=  dataset.map(__UpperCamelCase					).remove_columns(['''annotations''', '''document''', '''id''']					)
    print(__UpperCamelCase					)

    __UpperCamelCase     						=  dataset.filter(lambda lowercase_					: (len(x['''question''']					) + len(x['''context''']					)) < 4 * 40_96					)
    __UpperCamelCase     						=  short_validation_dataset.filter(lambda lowercase_					: x["category"] != "null"					)
    short_validation_dataset

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

    @jax.jit
    def forward(*lowercase_    ,   **lowercase_					):
        __UpperCamelCase						,					__UpperCamelCase						,					__UpperCamelCase     						=  model(*__UpperCamelCase    ,   **__UpperCamelCase					)
        return start_logits, end_logits, jnp.argmax(__UpperCamelCase    ,   axis=-1					)

    def evaluate(lowercase_					):
        # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
        __UpperCamelCase     						=  tokenizer(
            example['''question''']    ,   example['''context''']    ,   return_tensors='''np'''    ,   max_length=40_96    ,   padding='''max_length'''    ,   truncation=__UpperCamelCase    ,   )

        __UpperCamelCase						,					__UpperCamelCase						,					__UpperCamelCase     						=  forward(**__UpperCamelCase					)

        __UpperCamelCase     						=  CATEGORY_MAPPING[category.item()]

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

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

        __UpperCamelCase     						=  38 if predicted_category == '''short''' else 10_24
        __UpperCamelCase						,					__UpperCamelCase     						=  get_best_valid_start_end_idx(
            start_scores[0]    ,   end_scores[0]    ,   top_k=8    ,   max_size=__UpperCamelCase					)

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

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

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

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

        return example

    __UpperCamelCase     						=  short_validation_dataset.map(__UpperCamelCase					)

    __UpperCamelCase     						=  len(__UpperCamelCase					)
    __UpperCamelCase     						=  len(short_validation_dataset.filter(lambda lowercase_					: x["match"] == 1					)					)
    print('''EM score:'''    ,   (matched / total) * 1_00    ,   '''%'''					)


if __name__ == "__main__":
 main()












from __future__ import annotations

from enum import Enum, unique
from typing import TypeVar

# Create a generic variable that can be 'Enum', or any subclass.
_lowerCamelCase			 =						TypeVar('T', bound='Enum')



@unique
class 			a     (   _A			):

     '''simple docstring'''





     lowerCAmelCase				:       Optional[Any]   		=  80
     lowerCAmelCase				:       List[str]   		=  70
     lowerCAmelCase				:       Optional[Any]   		=  60
     lowerCAmelCase				:       Optional[Any]   		=  50
     lowerCAmelCase				:       List[Any]   		=  40
     lowerCAmelCase				:       Dict   		=  30
     lowerCAmelCase				:       List[Any]   		=  20
     lowerCAmelCase				:       List[Any]   		=  10



@unique
class 			a     (   _A			):

     '''simple docstring'''





     lowerCAmelCase				:       Dict   		=  24
     lowerCAmelCase				:       Union[str, Any]   		=  21
     lowerCAmelCase				:       Tuple   		=  18
     lowerCAmelCase				:       str   		=  15
     lowerCAmelCase				:       Dict   		=  12
     lowerCAmelCase				:       Dict   		=  9
     lowerCAmelCase				:       Tuple   		=  6
     lowerCAmelCase				:       Any   		=  3
     lowerCAmelCase				:       List[str]   		=  2
     lowerCAmelCase				:       Any   		=  1
     lowerCAmelCase				:       str   		=  -1
     lowerCAmelCase				:       Any   		=  -2
     lowerCAmelCase				:       List[str]   		=  -3
     lowerCAmelCase				:       Dict   		=  -6
     lowerCAmelCase				:       int   		=  -9
     lowerCAmelCase				:       Optional[Any]   		=  -12
     lowerCAmelCase				:       List[str]   		=  -15
     lowerCAmelCase				:       Any   		=  -18
     lowerCAmelCase				:       int   		=  -21
     lowerCAmelCase				:       Union[str, Any]   		=  -24
     @classmethod
     def   lowerCamelCase_   (	cls   :			type[T]				):
         return {unit.name: unit.value for unit in cls if unit.value > 0}




     @classmethod
     def   lowerCamelCase_   (	cls   :			type[T]				):
         return {unit.name: unit.value for unit in cls if unit.value < 0}




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						float      )  ->							str:
    UpperCAmelCase_							 =     SIUnit.get_positive() if value > 0 else SIUnit.get_negative()
    for name_prefix, value_prefix in prefixes.items():
        UpperCAmelCase_							 =     value / (10**value_prefix)
        if numerical_part > 1:
            return f'{numerical_part!s} {name_prefix}'
    return str(__UpperCamelCase      )




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						float      )  ->							str:
    for prefix in BinaryUnit:
        UpperCAmelCase_							 =     value / (2**prefix.value)
        if numerical_part > 1:
            return f'{numerical_part!s} {prefix.name}'
    return str(__UpperCamelCase      )


if __name__ == "__main__":
 import doctest

 doctest.testmod()





