





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

from transformers import BigBirdTokenizerFast


_UpperCamelCase       :				List[Any]           =			{0: """null""", 1: """short""", 2: """long""", 3: """yes""", 4: """no"""}
_UpperCamelCase       :				Tuple           =			set("""""".join(["""‘""", """’""", """´""", """`""", """.""", """,""", """-""", """\""""]))
def  __UpperCamelCase    (						snake_case							,			snake_case=0							,			snake_case=None      )    ->		List[str]:


   '''simple docstring'''







   return [" ".join(x.split(''' '''      )[begin:end]      ) for x in answers if len(x.split(''' '''      )      ) > 1]
def  __UpperCamelCase    (						snake_case							,			snake_case=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"
      __A    =  (
          given_answers + get_sub_answers(_A							,			begin=1      ) + get_sub_answers(_A							,			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(_A      )
def  __UpperCamelCase    (						snake_case							,			snake_case							,			snake_case=1							,			snake_case=1_0_0      )    ->		Union[str, Any]:


   '''simple docstring'''







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

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

   return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]
def  __UpperCamelCase    (						snake_case      )    ->		Union[str, Any]:


   '''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(_A      )      ) 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(_A      )      )

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

   __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  __UpperCamelCase    (						)    ->		Dict:


   '''simple docstring'''







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

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

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

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

   def evaluate(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_0_9_6							,			padding='''max_length'''							,			truncation=_A							,			)

      __A   ,		__A   ,		__A    =  forward(**_A      )

      __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    =  3_8 if predicted_category == '''short''' else 1_0_2_4
      __A   ,		__A    =  get_best_valid_start_end_idx(
          start_scores[0]							,			end_scores[0]							,			top_k=8							,			max_size=_A      )

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

      return example

   __A    =  short_validation_dataset.map(_A      )

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


if __name__ == "__main__":
    main()




from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase 	=   logging.get_logger(__name__)

lowerCamelCase 	=   {
    '''uw-madison/yoso-4096''': '''https://huggingface.co/uw-madison/yoso-4096/resolve/main/config.json''',
    # See all YOSO models at https://huggingface.co/models?filter=yoso
}







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   lowerCamelCase__					:      Tuple        =  'yoso'



   def __init__(							self,       lowercase_=50265,       lowercase_=768,       lowercase_=12,       lowercase_=12,       lowercase_=3072,       lowercase_="gelu",       lowercase_=0.1,       lowercase_=0.1,       lowercase_=4096,       lowercase_=1,       lowercase_=0.02,       lowercase_=1E-12,       lowercase_="absolute",       lowercase_=True,       lowercase_=9,       lowercase_=64,       lowercase_=None,       lowercase_=True,       lowercase_=True,       lowercase_=1,       lowercase_=0,       lowercase_=2,       **lowercase_,       )				->       List[str]:






     """simple docstring"""






     super().__init__(pad_token_id=lowercase_,       bos_token_id=lowercase_,       eos_token_id=lowercase_,       **lowercase_			)

     a__    =vocab_size
     a__    =max_position_embeddings
     a__    =hidden_size
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =intermediate_size
     a__    =hidden_act
     a__    =hidden_dropout_prob
     a__    =attention_probs_dropout_prob
     a__    =initializer_range
     a__    =type_vocab_size
     a__    =layer_norm_eps
     a__    =position_embedding_type
     a__    =use_expectation
     a__    =hash_code_len
     a__    =num_hash
     a__    =conv_window
     a__    =use_fast_hash
     a__    =lsh_backward
