

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

from transformers import BigBirdTokenizerFast


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




def  lowerCAmelCase(			SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_=0 ,	SCREAMING_SNAKE_CASE_=None )->	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 )->	Union[str, Any]:



    """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(SCREAMING_SNAKE_CASE_ ,	begin=1 ) + get_sub_answers(SCREAMING_SNAKE_CASE_ ,	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(SCREAMING_SNAKE_CASE_ )




def  lowerCAmelCase(			SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_=1 ,	SCREAMING_SNAKE_CASE_=1_0_0 )->	Optional[int]:



    """simple docstring"""






    UpperCamelCase_     ,					UpperCamelCase_      =	jax.lax.top_k(SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_ )
    UpperCamelCase_     ,					UpperCamelCase_      =	jax.lax.top_k(SCREAMING_SNAKE_CASE_ ,	SCREAMING_SNAKE_CASE_ )

    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(SCREAMING_SNAKE_CASE_ ).item()

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




def  lowerCAmelCase(			SCREAMING_SNAKE_CASE_ )->	int:



    """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(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_      =	[]
    for s in short_answers:
        short_targets.extend(s["text"] )
    UpperCamelCase_      =	list(set(SCREAMING_SNAKE_CASE_ ) )

    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(SCREAMING_SNAKE_CASE_ ) ) if not html[i]] )
        if new_answer not in long_targets:
            long_targets.append(SCREAMING_SNAKE_CASE_ )

    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  lowerCAmelCase(			)->	List[str]:



    """simple docstring"""






    UpperCamelCase_      =	load_from_disk("natural-questions-validation" )
    UpperCamelCase_      =	dataset.map(SCREAMING_SNAKE_CASE_ ).remove_columns(["annotations", "document", "id"] )
    print(SCREAMING_SNAKE_CASE_ )

    UpperCamelCase_      =	dataset.filter(lambda SCREAMING_SNAKE_CASE_ : (len(x["question"] ) + len(x["context"] )) < 4 * 4_0_9_6 )
    UpperCamelCase_      =	short_validation_dataset.filter(lambda SCREAMING_SNAKE_CASE_ : x["category"] != "null" )
    short_validation_dataset

    UpperCamelCase_      =	"vasudevgupta/flax-bigbird-natural-questions"
    UpperCamelCase_      =	FlaxBigBirdForNaturalQuestions.from_pretrained(SCREAMING_SNAKE_CASE_ )
    UpperCamelCase_      =	BigBirdTokenizerFast.from_pretrained(SCREAMING_SNAKE_CASE_ )

    @jax.jit
    def forward(*SCREAMING_SNAKE_CASE_ ,	**SCREAMING_SNAKE_CASE_ ):
        UpperCamelCase_     ,					UpperCamelCase_     ,					UpperCamelCase_      =	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
        UpperCamelCase_      =	tokenizer(
            example["question"] ,	example["context"] ,	return_tensors="np" ,	max_length=4_0_9_6 ,	padding="max_length" ,	truncation=SCREAMING_SNAKE_CASE_ ,	)

        UpperCamelCase_     ,					UpperCamelCase_     ,					UpperCamelCase_      =	forward(**SCREAMING_SNAKE_CASE_ )

        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_      =	3_8 if predicted_category == "short" else 1_0_2_4
        UpperCamelCase_     ,					UpperCamelCase_      =	get_best_valid_start_end_idx(
            start_scores[0] ,	end_scores[0] ,	top_k=8 ,	max_size=SCREAMING_SNAKE_CASE_ )

        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=SCREAMING_SNAKE_CASE_ )
        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(SCREAMING_SNAKE_CASE_ ) > 0

        return example

    UpperCamelCase_      =	short_validation_dataset.map(SCREAMING_SNAKE_CASE_ )

    UpperCamelCase_      =	len(SCREAMING_SNAKE_CASE_ )
    UpperCamelCase_      =	len(short_validation_dataset.filter(lambda SCREAMING_SNAKE_CASE_ : x["match"] == 1 ) )
    print("EM score:" ,	(matched / total) * 1_0_0 ,	"%" )


if __name__ == "__main__":
 main()



from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


SCREAMING_SNAKE_CASE		:Any				  =    logging.get_logger(__name__)

SCREAMING_SNAKE_CASE		:Dict				  =    {
    """facebook/blenderbot-3B""": """https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json""",
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}






class 	__magic_name__				(       snake_case							):
     UpperCamelCase_  :str								=			"""blenderbot"""
     UpperCamelCase_  :Union[str, Any]								=			["""past_key_values"""]
     UpperCamelCase_  :Any								=			{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






     def __init__(							self		,	_lowercase=8_008		,	_lowercase=128		,	_lowercase=2		,	_lowercase=10_240		,	_lowercase=32		,	_lowercase=24		,	_lowercase=10_240		,	_lowercase=32		,	_lowercase=0.0		,	_lowercase=0.0		,	_lowercase=True		,	_lowercase=True		,	_lowercase="gelu"		,	_lowercase=2_560		,	_lowercase=0.1		,	_lowercase=0.0		,	_lowercase=0.0		,	_lowercase=0.02		,	_lowercase=1		,	_lowercase=False		,	_lowercase=0		,	_lowercase=1		,	_lowercase=2		,	_lowercase=3		,	_lowercase=2		,	**_lowercase		,	)->    List[str]:
         UpperCamelCase_      =	vocab_size
         UpperCamelCase_      =	max_position_embeddings
         UpperCamelCase_      =	d_model
         UpperCamelCase_      =	encoder_ffn_dim
         UpperCamelCase_      =	encoder_layers
         UpperCamelCase_      =	encoder_attention_heads
         UpperCamelCase_      =	decoder_ffn_dim
         UpperCamelCase_      =	decoder_layers
         UpperCamelCase_      =	decoder_attention_heads
         UpperCamelCase_      =	dropout
         UpperCamelCase_      =	attention_dropout
         UpperCamelCase_      =	activation_dropout
         UpperCamelCase_      =	activation_function
         UpperCamelCase_      =	init_std
         UpperCamelCase_      =	encoder_layerdrop
         UpperCamelCase_      =	decoder_layerdrop
         UpperCamelCase_      =	use_cache
         UpperCamelCase_      =	encoder_layers
         UpperCamelCase_      =	scale_embedding  # scale factor will be sqrt(d_model) if True

         super().__init__(
             pad_token_id=_lowercase		,	bos_token_id=_lowercase		,	eos_token_id=_lowercase		,	is_encoder_decoder=_lowercase		,	decoder_start_token_id=_lowercase		,	encoder_no_repeat_ngram_size=_lowercase		,	forced_eos_token_id=_lowercase		,	**_lowercase		,	)






class 	__magic_name__				(       snake_case							):






     @property
     def   UpperCAmelCase_							(							self  )->    Mapping[str, Mapping[int, str]]:
         if self.task in ["default", "seq2seq-lm"]:
             UpperCamelCase_      =	OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                 ]  )
             if self.use_past:
                 UpperCamelCase_      =	{0: "batch"}
                 UpperCamelCase_      =	{0: "batch", 1: "past_decoder_sequence + sequence"}
             else:
                 UpperCamelCase_      =	{0: "batch", 1: "decoder_sequence"}
                 UpperCamelCase_      =	{0: "batch", 1: "decoder_sequence"}
             if self.use_past:
                 self.fill_with_past_key_values_(_lowercase		,	direction="inputs"  )
         elif self.task == "causal-lm":
             UpperCamelCase_      =	OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                 ]  )
             if self.use_past:
                 UpperCamelCase_     ,					UpperCamelCase_      =	self.num_layers
                 for i in range(_lowercase  ):
                     UpperCamelCase_      =	{0: "batch", 2: "past_sequence + sequence"}
                     UpperCamelCase_      =	{0: "batch", 2: "past_sequence + sequence"}
         else:
             UpperCamelCase_      =	OrderedDict(
                 [
                     ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                     ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                     ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
                     ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
                 ]  )

         return common_inputs






     @property
     # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
     def   UpperCAmelCase_							(							self  )->    Mapping[str, Mapping[int, str]]:
         if self.task in ["default", "seq2seq-lm"]:
             UpperCamelCase_      =	super().outputs
         else:
             UpperCamelCase_      =	super(_lowercase		,	self  ).outputs
             if self.use_past:
                 UpperCamelCase_     ,					UpperCamelCase_      =	self.num_layers
                 for i in range(_lowercase  ):
                     UpperCamelCase_      =	{0: "batch", 2: "past_sequence + sequence"}
                     UpperCamelCase_      =	{0: "batch", 2: "past_sequence + sequence"}
         return common_outputs






     def   UpperCAmelCase_							(							self		,	_lowercase		,	_lowercase = -1		,	_lowercase = -1		,	_lowercase = False		,	_lowercase = None		,	)->    Mapping[str, Any]:
         UpperCamelCase_      =	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             _lowercase		,	_lowercase		,	_lowercase		,	_lowercase		,	_lowercase  )
         # Generate decoder inputs
         UpperCamelCase_      =	seq_length if not self.use_past else 1
         UpperCamelCase_      =	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             _lowercase		,	_lowercase		,	_lowercase		,	_lowercase		,	_lowercase  )
         UpperCamelCase_      =	{F"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
         UpperCamelCase_      =	dict(**_lowercase		,	**_lowercase  )

         if self.use_past:
             if not is_torch_available():
                 raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."  )
             else:
                 import torch
             UpperCamelCase_     ,					UpperCamelCase_      =	common_inputs["input_ids"].shape
             UpperCamelCase_      =	common_inputs["decoder_input_ids"].shape[1]
             UpperCamelCase_     ,					UpperCamelCase_      =	self.num_attention_heads
             UpperCamelCase_      =	(
                 batch,
                 num_encoder_attention_heads,
                 encoder_seq_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             UpperCamelCase_      =	decoder_seq_length
             UpperCamelCase_      =	(
                 batch,
                 num_decoder_attention_heads,
                 decoder_past_length,
                 self._config.hidden_size // num_decoder_attention_heads,
             )
             UpperCamelCase_      =	torch.cat(
                 [common_inputs["decoder_attention_mask"], torch.ones(_lowercase		,	_lowercase  )]		,	dim=1  )
             UpperCamelCase_      =	[]
             UpperCamelCase_     ,					UpperCamelCase_      =	self.num_layers

             for _ in range(_lowercase  ):
                 common_inputs["past_key_values"].append(
                     (
                         torch.zeros(_lowercase  ),
                         torch.zeros(_lowercase  ),
                         torch.zeros(_lowercase  ),
                         torch.zeros(_lowercase  ),
                     )  )
         return common_inputs






     def   UpperCAmelCase_							(							self		,	_lowercase		,	_lowercase = -1		,	_lowercase = -1		,	_lowercase = False		,	_lowercase = None		,	)->    Mapping[str, Any]:
         UpperCamelCase_      =	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             _lowercase		,	_lowercase		,	_lowercase		,	_lowercase		,	_lowercase  )

         if self.use_past:
             if not is_torch_available():
                 raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."  )
             else:
                 import torch
             UpperCamelCase_     ,					UpperCamelCase_      =	common_inputs["input_ids"].shape
             UpperCamelCase_      =	seqlen
             UpperCamelCase_     ,					UpperCamelCase_      =	self.num_layers
             UpperCamelCase_     ,					UpperCamelCase_      =	self.num_attention_heads
             UpperCamelCase_      =	(
                 batch,
                 num_encoder_attention_heads,
                 past_key_values_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             UpperCamelCase_      =	common_inputs["attention_mask"].dtype
             UpperCamelCase_      =	torch.cat(
                 [common_inputs["attention_mask"], torch.ones(_lowercase		,	_lowercase		,	dtype=_lowercase  )]		,	dim=1  )
             UpperCamelCase_      =	[
                 (torch.zeros(_lowercase  ), torch.zeros(_lowercase  )) for _ in range(_lowercase  )
             ]
         return common_inputs






     def   UpperCAmelCase_							(							self		,	_lowercase		,	_lowercase = -1		,	_lowercase = -1		,	_lowercase = False		,	_lowercase = None		,	)->    Mapping[str, Any]:
         # Copied from OnnxConfig.generate_dummy_inputs
         # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
         # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
         UpperCamelCase_      =	compute_effective_axis_dimension(
             _lowercase		,	fixed_dimension=OnnxConfig.default_fixed_batch		,	num_token_to_add=0  )

         # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
         UpperCamelCase_      =	tokenizer.num_special_tokens_to_add(_lowercase  )
         UpperCamelCase_      =	compute_effective_axis_dimension(
             _lowercase		,	fixed_dimension=OnnxConfig.default_fixed_sequence		,	num_token_to_add=_lowercase  )

         # Generate dummy inputs according to compute batch and sequence
         UpperCamelCase_      =	[" ".join([tokenizer.unk_token]  ) * seq_length] * batch_size
         UpperCamelCase_      =	dict(tokenizer(_lowercase		,	return_tensors=_lowercase  )  )
         return common_inputs






     def   UpperCAmelCase_							(							self		,	_lowercase		,	_lowercase = -1		,	_lowercase = -1		,	_lowercase = False		,	_lowercase = None		,	)->    Mapping[str, Any]:
         if self.task in ["default", "seq2seq-lm"]:
             UpperCamelCase_      =	self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                 _lowercase		,	batch_size=_lowercase		,	seq_length=_lowercase		,	is_pair=_lowercase		,	framework=_lowercase  )

         elif self.task == "causal-lm":
             UpperCamelCase_      =	self._generate_dummy_inputs_for_causal_lm(
                 _lowercase		,	batch_size=_lowercase		,	seq_length=_lowercase		,	is_pair=_lowercase		,	framework=_lowercase  )
         else:
             UpperCamelCase_      =	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                 _lowercase		,	batch_size=_lowercase		,	seq_length=_lowercase		,	is_pair=_lowercase		,	framework=_lowercase  )

         return common_inputs






     def   UpperCAmelCase_							(							self		,	_lowercase		,	_lowercase		,	_lowercase		,	_lowercase  )->    List[Any]:
         if self.task in ["default", "seq2seq-lm"]:
             UpperCamelCase_      =	super()._flatten_past_key_values_(_lowercase		,	_lowercase		,	_lowercase		,	_lowercase  )
         else:
             UpperCamelCase_      =	super(_lowercase		,	self  )._flatten_past_key_values_(
                 _lowercase		,	_lowercase		,	_lowercase		,	_lowercase  )






     def   UpperCAmelCase_							(							self		,	_lowercase		,	_lowercase  )->    Union[str, Any]:
         if direction not in ["inputs", "outputs"]:
             raise ValueError(F"direction must either be \"inputs\" or \"outputs\", but {direction} was given"  )

         UpperCamelCase_      =	"past_key_values" if direction == "inputs" else "present"
         UpperCamelCase_     ,					UpperCamelCase_      =	self.num_layers

         UpperCamelCase_      =	"past_encoder_sequence"
         UpperCamelCase_      =	"past_decoder_sequence" if direction == "inputs" else "past_decoder_sequence + sequence"

         for i in range(_lowercase  ):
             UpperCamelCase_      =	{0: "batch", 2: decoder_sequence}
             UpperCamelCase_      =	{0: "batch", 2: decoder_sequence}
             UpperCamelCase_      =	{0: "batch", 2: encoder_sequence}
             UpperCamelCase_      =	{0: "batch", 2: encoder_sequence}
