





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       : List[Any]						    = set("".join(["‘", "’", "´", "`", ".", ",", "-", "\""]))





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE=0       ,    _SCREAMING_SNAKE_CASE=None							)       ->   List[Any]:




    """simple docstring"""




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





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE=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"
        _A    =       (
            given_answers + get_sub_answers(_SCREAMING_SNAKE_CASE       ,    begin=1							) + get_sub_answers(_SCREAMING_SNAKE_CASE       ,    end=-1							)
        )
    _A    =       []
    for answer in given_answers:
        _A    =       answer.replace('_'       ,    ' '							).lower()
        _A    =       ''.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 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE=1       ,    _SCREAMING_SNAKE_CASE=100							)       ->   List[str]:




    """simple docstring"""




    _A, _A    =       jax.lax.top_k(_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)
    _A, _A    =       jax.lax.top_k(_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)

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

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





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE							)       ->   Dict:




    """simple docstring"""




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

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

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

    _A    =       []
    for s in long_answers:
        if s["start_token"] == -1:
            continue
        _A    =       context[s['start_token'] : s['end_token']]
        _A    =       is_html[s['start_token'] : s['end_token']]
        _A    =       ' '.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							)

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




    """simple docstring"""




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

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

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

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

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

        _A, _A, _A    =       forward(**_SCREAMING_SNAKE_CASE							)

        _A    =       CATEGORY_MAPPING[category.item()]

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

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

        _A    =       38 if predicted_category == 'short' else 1_024
        _A, _A    =       get_best_valid_start_end_idx(
            start_scores[0]       ,    end_scores[0]       ,    top_k=8       ,    max_size=_SCREAMING_SNAKE_CASE							)

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

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

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

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

        return example

    _A    =       short_validation_dataset.map(_SCREAMING_SNAKE_CASE							)

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


if __name__ == "__main__":
 main()












import numpy as np





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)       ->   List[Any]:




    """simple docstring"""




    _A    =       np.array(_SCREAMING_SNAKE_CASE							)
    _A    =       np.array(_SCREAMING_SNAKE_CASE							)

    _A    =       abs(predict - actual							)
    _A    =       difference.mean()

    return score





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)       ->   Optional[int]:




    """simple docstring"""




    _A    =       np.array(_SCREAMING_SNAKE_CASE							)
    _A    =       np.array(_SCREAMING_SNAKE_CASE							)

    _A    =       predict - actual
    _A    =       np.square(_SCREAMING_SNAKE_CASE							)

    _A    =       square_diff.mean()
    return score





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)       ->   Tuple:




    """simple docstring"""




    _A    =       np.array(_SCREAMING_SNAKE_CASE							)
    _A    =       np.array(_SCREAMING_SNAKE_CASE							)

    _A    =       predict - actual
    _A    =       np.square(_SCREAMING_SNAKE_CASE							)
    _A    =       square_diff.mean()
    _A    =       np.sqrt(_SCREAMING_SNAKE_CASE							)
    return score





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)       ->   List[Any]:




    """simple docstring"""




    _A    =       np.array(_SCREAMING_SNAKE_CASE							)
    _A    =       np.array(_SCREAMING_SNAKE_CASE							)

    _A    =       np.log(predict + 1							)
    _A    =       np.log(actual + 1							)

    _A    =       log_predict - log_actual
    _A    =       np.square(_SCREAMING_SNAKE_CASE							)
    _A    =       square_diff.mean()

    _A    =       np.sqrt(_SCREAMING_SNAKE_CASE							)

    return score





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)       ->   Optional[int]:




    """simple docstring"""




    _A    =       np.array(_SCREAMING_SNAKE_CASE							)
    _A    =       np.array(_SCREAMING_SNAKE_CASE							)

    _A    =       predict - actual
    _A    =       np.sum(_SCREAMING_SNAKE_CASE							) / len(_SCREAMING_SNAKE_CASE							)
    _A    =       np.sum(_SCREAMING_SNAKE_CASE							) / len(_SCREAMING_SNAKE_CASE							)
    # print(numerator, denumerator)
    _A    =       float(_SCREAMING_SNAKE_CASE							) / denumerator * 100

    return score





def 							__lowerCAmelCase(					_SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE							)       ->   Dict:




    """simple docstring"""




    return np.mean(np.array(_SCREAMING_SNAKE_CASE							) == np.array(_SCREAMING_SNAKE_CASE							)							)





