


'''simple docstring'''

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

from transformers import BigBirdTokenizerFast


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

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






       """simple docstring"""







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

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






       """simple docstring"""







       lowercase__      ,lowercase__				      =			jax.lax.top_k(A  ,			A	)
       lowercase__      ,lowercase__				      =			jax.lax.top_k(A  ,			A	)

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

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

def _SCREAMING_SNAKE_CASE				(A	)	->	Any:






       """simple docstring"""







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

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

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

       lowercase__				      =			[]
       for s in long_answers:
              if s["start_token"] == -1:
                     continue
              lowercase__				      =			context[s['''start_token'''] : s['''end_token''']]
              lowercase__				      =			is_html[s['''start_token'''] : s['''end_token''']]
              lowercase__				      =			''' '''.join([answer[i] for i in range(len(A	)	) if not html[i]]	)
              if new_answer not in long_targets:
                     long_targets.append(A	)

       lowercase__				      =			'''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 _SCREAMING_SNAKE_CASE				()	->	str:






       """simple docstring"""







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

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

       lowercase__				      =			'''vasudevgupta/flax-bigbird-natural-questions'''
       lowercase__				      =			FlaxBigBirdForNaturalQuestions.from_pretrained(A	)
       lowercase__				      =			BigBirdTokenizerFast.from_pretrained(A	)

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

              lowercase__      ,lowercase__      ,lowercase__				      =			forward(**A	)

              lowercase__				      =			CATEGORY_MAPPING[category.item()]

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

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

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

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

              lowercase__				      =			expand_to_aliases(example['''targets''']  ,			make_sub_answers=A	)
              lowercase__				      =			expand_to_aliases(example['''output''']	)

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

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

              return example

       lowercase__				      =			short_validation_dataset.map(A	)

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


if __name__ == "__main__":
 main()








'''simple docstring'''

def _SCREAMING_SNAKE_CASE				(A	)	->	list:






       """simple docstring"""







       lowercase__				      =			len(A	)

       # Set the gap value to be decreased by a factor of 1.3
       # after each iteration
       lowercase__				      =			1.3

       # Continue sorting until the gap is 1
       while gap > 1:
              # Decrease the gap value
              lowercase__				      =			int(gap / shrink	)

              # Sort the elements using insertion sort
              for i in range(A  ,			len(A	)	):
                     lowercase__				      =			collection[i]
                     lowercase__				      =			i
                     while j >= gap and collection[j - gap] > temp:
                            lowercase__				      =			collection[j - gap]
                            j -= gap
                     lowercase__				      =			temp

       return collection


if __name__ == "__main__":
 import doctest

 doctest.testmod()




