






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

from transformers import BigBirdTokenizerFast


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


def   _A       (    SCREAMING_SNAKE_CASE__  :     Optional[int]					,  SCREAMING_SNAKE_CASE__  :     Union[str, Any]=0					,  SCREAMING_SNAKE_CASE__  :     List[str]=None	):
    return [" ".join(x.split(''' '''	)[begin:end]	) for x in answers if len(x.split(''' '''	)	) > 1]


def   _A       (    SCREAMING_SNAKE_CASE__  :     List[Any]					,  SCREAMING_SNAKE_CASE__  :     Tuple=False	):
    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	:Union[str, Any]     							=       (
            given_answers + get_sub_answers(SCREAMING_SNAKE_CASE__					,  begin=1	) + get_sub_answers(SCREAMING_SNAKE_CASE__					,  end=-1	)
        )
    UpperCamelCase	:Union[str, Any]     							=       []
    for answer in given_answers:
        UpperCamelCase	:List[Any]     							=       answer.replace('''_'''					,  ''' '''	).lower()
        UpperCamelCase	:str     							=       ''''''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ''' ''' for c in alias	)
        answers.append(''' '''.join(alias.split()	).strip()	)
    return set(SCREAMING_SNAKE_CASE__	)


def   _A       (    SCREAMING_SNAKE_CASE__  :     List[str]					,  SCREAMING_SNAKE_CASE__  :     int					,  SCREAMING_SNAKE_CASE__  :     Dict=1					,  SCREAMING_SNAKE_CASE__  :     Tuple=100	):
    UpperCamelCase							,       UpperCamelCase	:Optional[Any]     							=       jax.lax.top_k(SCREAMING_SNAKE_CASE__					,  SCREAMING_SNAKE_CASE__	)
    UpperCamelCase							,       UpperCamelCase	:List[Any]     							=       jax.lax.top_k(SCREAMING_SNAKE_CASE__					,  SCREAMING_SNAKE_CASE__	)

    UpperCamelCase	:Any     							=       best_end_idx[:, None] - best_start_idx[None, :]
    UpperCamelCase	:List[str]     							=       jnp.logical_or(widths < 0					,  widths > max_size	)
    UpperCamelCase	:List[Any]     							=       (best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
    UpperCamelCase	:Optional[int]     							=       jnp.argmax(SCREAMING_SNAKE_CASE__	).item()

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


def   _A       (    SCREAMING_SNAKE_CASE__  :     List[str]	):
    UpperCamelCase	:List[str]     							=       sample['''question''']['''text''']
    UpperCamelCase	:str     							=       sample['''document''']['''tokens''']['''token''']
    UpperCamelCase	:Union[str, Any]     							=       sample['''document''']['''tokens''']['''is_html''']
    UpperCamelCase	:Any     							=       sample['''annotations''']['''long_answer''']
    UpperCamelCase	:Optional[int]     							=       sample['''annotations''']['''short_answers''']

    UpperCamelCase	:List[Any]     							=       ''' '''.join([context[i] for i in range(len(SCREAMING_SNAKE_CASE__	)	) 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	:List[str]     							=       []
    for s in short_answers:
        short_targets.extend(s['''text''']	)
    UpperCamelCase	:Optional[Any]     							=       list(set(SCREAMING_SNAKE_CASE__	)	)

    UpperCamelCase	:List[str]     							=       []
    for s in long_answers:
        if s["start_token"] == -1:
            continue
        UpperCamelCase	:List[str]     							=       context[s['''start_token'''] : s['''end_token''']]
        UpperCamelCase	:Dict     							=       is_html[s['''start_token'''] : s['''end_token''']]
        UpperCamelCase	:Union[str, Any]     							=       ''' '''.join([answer[i] for i in range(len(SCREAMING_SNAKE_CASE__	)	) if not html[i]]	)
        if new_answer not in long_targets:
            long_targets.append(SCREAMING_SNAKE_CASE__	)

    UpperCamelCase	:Union[str, Any]     							=       '''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       (    ):
    UpperCamelCase	:int     							=       load_from_disk('''natural-questions-validation'''	)
    UpperCamelCase	:Optional[int]     							=       dataset.map(SCREAMING_SNAKE_CASE__	).remove_columns(['''annotations''', '''document''', '''id''']	)
    print(SCREAMING_SNAKE_CASE__	)

    UpperCamelCase	:List[str]     							=       dataset.filter(lambda SCREAMING_SNAKE_CASE__	: (len(x['''question''']	) + len(x['''context''']	)) < 4 * 4096	)
    UpperCamelCase	:Dict     							=       short_validation_dataset.filter(lambda SCREAMING_SNAKE_CASE__	: x["category"] != "null"	)
    short_validation_dataset

    UpperCamelCase	:Union[str, Any]     							=       '''vasudevgupta/flax-bigbird-natural-questions'''
    UpperCamelCase	:Union[str, Any]     							=       FlaxBigBirdForNaturalQuestions.from_pretrained(SCREAMING_SNAKE_CASE__	)
    UpperCamelCase	:Dict     							=       BigBirdTokenizerFast.from_pretrained(SCREAMING_SNAKE_CASE__	)

    @jax.jit
    def forward(*SCREAMING_SNAKE_CASE__  :     Optional[Any]					,  **SCREAMING_SNAKE_CASE__  :     int	):
        UpperCamelCase							,       UpperCamelCase							,       UpperCamelCase	:int     							=       model(*SCREAMING_SNAKE_CASE__					,  **SCREAMING_SNAKE_CASE__	)
        return start_logits, end_logits, jnp.argmax(SCREAMING_SNAKE_CASE__					,  axis=-1	)

    def evaluate(SCREAMING_SNAKE_CASE__  :     str	):
        # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
        UpperCamelCase	:List[Any]     							=       tokenizer(
            example['''question''']					,  example['''context''']					,  return_tensors='''np'''					,  max_length=4096					,  padding='''max_length'''					,  truncation=SCREAMING_SNAKE_CASE__					,  )

        UpperCamelCase							,       UpperCamelCase							,       UpperCamelCase	:Any     							=       forward(**SCREAMING_SNAKE_CASE__	)

        UpperCamelCase	:Tuple     							=       CATEGORY_MAPPING[category.item()]

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

        if predicted_category in ["yes", "no", "null"]:
            UpperCamelCase	:List[str]     							=       [predicted_category]
            UpperCamelCase	:Optional[Any]     							=       example['''output'''] == example['''targets''']
            UpperCamelCase	:Union[str, Any]     							=       predicted_category != '''null'''
            return example

        UpperCamelCase	:Dict     							=       38 if predicted_category == '''short''' else 1024
        UpperCamelCase							,       UpperCamelCase	:Optional[Any]     							=       get_best_valid_start_end_idx(
            start_scores[0]					,  end_scores[0]					,  top_k=8					,  max_size=SCREAMING_SNAKE_CASE__	)

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

        UpperCamelCase	:Optional[int]     							=       expand_to_aliases(example['''targets''']					,  make_sub_answers=SCREAMING_SNAKE_CASE__	)
        UpperCamelCase	:int     							=       expand_to_aliases(example['''output''']	)

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

        # if there is a common element, it's a exact match
        UpperCamelCase	:Dict     							=       len(list(answers & predictions	)	) > 0
        UpperCamelCase	:Union[str, Any]     							=       predicted_category != '''null''' and len(SCREAMING_SNAKE_CASE__	) > 0

        return example

    UpperCamelCase	:Union[str, Any]     							=       short_validation_dataset.map(SCREAMING_SNAKE_CASE__	)

    UpperCamelCase	:Tuple     							=       len(SCREAMING_SNAKE_CASE__	)
    UpperCamelCase	:Dict     							=       len(short_validation_dataset.filter(lambda SCREAMING_SNAKE_CASE__	: x["match"] == 1	)	)
    print('''EM score:'''					,  (matched / total) * 100					,  '''%'''	)


if __name__ == "__main__":
 main()













def   _A       (    SCREAMING_SNAKE_CASE__  :     int	):
    UpperCamelCase	:Optional[Any]     							=       9

    while num > 0:
        if num % 10 != digit:
            return False
        num //= 100
        digit -= 1

    return True


def   _A       (    ):
    UpperCamelCase	:List[str]     							=       138902663

    while not is_square_form(num * num	):
        if num % 10 == 3:
            num -= 6  # (3 - 6) % 10 = 7
        else:
            num -= 4  # (7 - 4) % 10 = 3

    return num * 10


if __name__ == "__main__":
 print(f'''{solution() = }''')





