



'''simple docstring'''



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

from transformers import BigBirdTokenizerFast


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







def _lowercase   (  __A							,__A=0							,__A=None       ):



 '''simple docstring'''



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







def _lowercase   (  __A							,__A=False       ):



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







def _lowercase   (  __A							,__A							,__A=1							,__A=100       ):



 '''simple docstring'''



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

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

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







def _lowercase   (  __A       ):



 '''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(__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",
   }

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

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

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



 '''simple docstring'''



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

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

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

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

 def evaluate(__A       ):
  # 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_096							,padding="""max_length"""							,truncation=__A							,)

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

  __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 1_024
  __UpperCamelCase      ,		__UpperCamelCase				      =  get_best_valid_start_end_idx(
      start_scores[0]							,end_scores[0]							,top_k=8							,max_size=__A       )

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

  return example

 __UpperCamelCase				      =  short_validation_dataset.map(__A       )

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


if __name__ == "__main__":
    main()










'''simple docstring'''



a__						:			Any										=			{
    '0000': '0',
    '0001': '1',
    '0010': '2',
    '0011': '3',
    '0100': '4',
    '0101': '5',
    '0110': '6',
    '0111': '7',
    '1000': '8',
    '1001': '9',
    '1010': 'a',
    '1011': 'b',
    '1100': 'c',
    '1101': 'd',
    '1110': 'e',
    '1111': 'f',
}







def _lowercase   (  __A       ):



 '''simple docstring'''



 __UpperCamelCase				      =  str(__A       ).strip()

 # Exceptions
 if not binary_str:
  raise ValueError("""Empty string was passed to the function"""       )
 __UpperCamelCase				      =  binary_str[0] == """-"""
 __UpperCamelCase				      =  binary_str[1:] if is_negative else binary_str
 if not all(char in """01""" for char in binary_str       ):
  raise ValueError("""Non-binary value was passed to the function"""       )

 __UpperCamelCase				      =  (
     """0""" * (4 * (divmod(len(__A       )							,4       )[0] + 1) - len(__A       )) + binary_str
 )

 __UpperCamelCase				      =  []
 for x in range(0							,len(__A       )							,4       ):
  hexadecimal.append(BITS_TO_HEX[binary_str[x : x + 4]]       )
 __UpperCamelCase				      =  """0x""" + """""".join(__A       )

 return "-" + hexadecimal_str if is_negative else hexadecimal_str


if __name__ == "__main__":
 from doctest import testmod

 testmod()





