





'''simple docstring'''

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

from transformers import BigBirdTokenizerFast


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







def       a__       (       lowercase      :  List[str],      lowercase      :  str=0,      lowercase      :  Dict=None   )	->	Dict:





					"""simple docstring"""

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







def       a__       (       lowercase      :  Tuple,      lowercase      :  Optional[Any]=False   )	->	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"
										_UpperCamelCase       					= (
										    given_answers + get_sub_answers(_UpperCAmelCase,      begin=1   ) + get_sub_answers(_UpperCAmelCase,      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(_UpperCAmelCase   )







def       a__       (       lowercase      :  Tuple,      lowercase      :  Dict,      lowercase      :  List[str]=1,      lowercase      :  List[str]=100   )	->	Tuple:





					"""simple docstring"""

					_UpperCamelCase       					= jax.lax.top_k(_UpperCAmelCase,      _UpperCAmelCase   )
					_UpperCamelCase       					= jax.lax.top_k(_UpperCAmelCase,      _UpperCAmelCase   )

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

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







def       a__       (       lowercase      :  Any   )	->	Optional[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(_UpperCAmelCase   )   ) 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(_UpperCAmelCase   )   )

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

					_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       a__       (       )	->	Union[str, Any]:





					"""simple docstring"""

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

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

					_UpperCamelCase       					= 'vasudevgupta/flax-bigbird-natural-questions'
					_UpperCamelCase       					= FlaxBigBirdForNaturalQuestions.from_pretrained(_UpperCAmelCase   )
					_UpperCamelCase       					= BigBirdTokenizerFast.from_pretrained(_UpperCAmelCase   )

					@jax.jit
					def forward(*lowercase      :  str,      **lowercase      :  Optional[Any]   ):
										_UpperCamelCase       					= model(*_UpperCAmelCase,      **_UpperCAmelCase   )
										return start_logits, end_logits, jnp.argmax(_UpperCAmelCase,      axis=-1   )

					def evaluate(lowercase      :  str   ):
										# 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=4096,      padding='''max_length''',      truncation=_UpperCAmelCase,      )

										_UpperCamelCase       					= forward(**_UpperCAmelCase   )

										_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 1024
										_UpperCamelCase       					= get_best_valid_start_end_idx(
										    start_scores[0],      end_scores[0],      top_k=8,      max_size=_UpperCAmelCase   )

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

										return example

					_UpperCamelCase       					= short_validation_dataset.map(_UpperCAmelCase   )

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


if __name__ == "__main__":
			main()










'''simple docstring'''

def       a__       (       lowercase      :  int   )	->	str:





     """simple docstring"""

     if not isinstance(lowercase,      lowercase   ):
          _UpperCamelCase       					= (
              '''operation can not be conducted on a object of type '''
              F"""{type(lowercase   ).__name__}"""
          )
          raise TypeError(lowercase   )
     _UpperCamelCase       					= ''''''
     for _ in range(0,      32   ):
          bit_string += str(number % 2   )
          _UpperCamelCase       					= number >> 1
     return bit_string







def       a__       (       lowercase      :  int   )	->	str:





     """simple docstring"""

     if number < 0:
          raise ValueError('''the value of input must be positive'''   )
     elif isinstance(lowercase,      lowercase   ):
          raise TypeError('''Input value must be a \'int\' type'''   )
     elif isinstance(lowercase,      lowercase   ):
          raise TypeError('''\'<\' not supported between instances of \'str\' and \'int\''''   )
     _UpperCamelCase       					= 0
     # iterator over [1 to 32],since we are dealing with 32 bit integer
     for _ in range(1,      33   ):
          # left shift the bits by unity
          _UpperCamelCase       					= result << 1
          # get the end bit
          _UpperCamelCase       					= number % 2
          # right shift the bits by unity
          _UpperCamelCase       					= number >> 1
          # add that bit to our ans
          _UpperCamelCase       					= result | end_bit
     return get_reverse_bit_string(lowercase   )


if __name__ == "__main__":
   import doctest

   doctest.testmod()



