


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

from transformers import BigBirdTokenizerFast


UpperCAmelCase__        =				{0: '''null''', 1: '''short''', 2: '''long''', 3: '''yes''', 4: '''no'''}
UpperCAmelCase__        =				set(''''''.join(['''‘''', '''’''', '''´''', '''`''', '''.''', ''',''', '''-''', '''\"''']))
def 							UpperCAmelCase_   (				__snake_case	,		__snake_case=0	,		__snake_case=None		)	->     Optional[int]:







 """simple docstring"""


 return [" ".join(x.split(''' '''		)[begin:end]		) for x in answers if len(x.split(''' '''		)		) > 1]
def 							UpperCAmelCase_   (				__snake_case	,		__snake_case=False		)	->     Tuple:







 """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(UpperCAmelCase_	,		begin=1		) + get_sub_answers(UpperCAmelCase_	,		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(UpperCAmelCase_		)
def 							UpperCAmelCase_   (				__snake_case	,		__snake_case	,		__snake_case=1	,		__snake_case=100		)	->     str:







 """simple docstring"""


 _lowercase				 =jax.lax.top_k(UpperCAmelCase_	,		UpperCAmelCase_		)
 _lowercase				 =jax.lax.top_k(UpperCAmelCase_	,		UpperCAmelCase_		)

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

 return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]
def 							UpperCAmelCase_   (				__snake_case		)	->     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(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",
   }

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

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

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







 """simple docstring"""


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

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

 _lowercase				 ='vasudevgupta/flax-bigbird-natural-questions'
 _lowercase				 =FlaxBigBirdForNaturalQuestions.from_pretrained(UpperCAmelCase_		)
 _lowercase				 =BigBirdTokenizerFast.from_pretrained(UpperCAmelCase_		)

 @jax.jit
 def forward(*__snake_case	,		**__snake_case		):
  _lowercase				 =model(*UpperCAmelCase_	,		**UpperCAmelCase_		)
  return start_logits, end_logits, jnp.argmax(UpperCAmelCase_	,		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
  _lowercase				 =tokenizer(
      example['''question''']	,		example['''context''']	,		return_tensors='''np'''	,		max_length=4096	,		padding='''max_length'''	,		truncation=UpperCAmelCase_	,		)

  _lowercase				 =forward(**UpperCAmelCase_		)

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

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

  return example

 _lowercase				 =short_validation_dataset.map(UpperCAmelCase_		)

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


if __name__ == "__main__":
  main()




import unittest

from cer import CER


lowerCAmelCase__             =		CER()







class 						lowercase     (					unittest.TestCase				):


		"""simple docstring"""




		def 		A__       (					self):
			_UpperCamelCase :		Any								=    ['White House']
			_UpperCamelCase :		Tuple								=    ['white house']
			# S = 2, D = 0, I = 0, N = 11, CER = 2 / 11
			_UpperCamelCase :		Tuple								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.1_8_1_8_1_8_1_8_1_8) < 1e-6)


		def 		A__       (					self):
			_UpperCamelCase :		Union[str, Any]								=    ['were wolf']
			_UpperCamelCase :		List[Any]								=    ['werewolf']
			# S = 0, D = 0, I = 1, N = 9, CER = 1 / 9
			_UpperCamelCase :		Dict								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.1_1_1_1_1_1_1) < 1e-6)

			_UpperCamelCase :		int								=    ['werewolf']
			_UpperCamelCase :		Optional[Any]								=    ['weae     wolf']
			# S = 1, D = 1, I = 0, N = 8, CER = 0.25
			_UpperCamelCase :		str								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.2_5) < 1e-6)

			# consecutive whitespaces case 1
			_UpperCamelCase :		Tuple								=    ['were wolf']
			_UpperCamelCase :		Dict								=    ['were               wolf']
			# S = 0, D = 0, I = 0, N = 9, CER = 0
			_UpperCamelCase :		List[str]								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.0) < 1e-6)

			# consecutive whitespaces case 2
			_UpperCamelCase :		Dict								=    ['were   wolf']
			_UpperCamelCase :		Dict								=    ['were               wolf']
			# S = 0, D = 0, I = 0, N = 9, CER = 0
			_UpperCamelCase :		Dict								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.0) < 1e-6)


		def 		A__       (					self):
			_UpperCamelCase :		Tuple								=    ['werewolf']
			_UpperCamelCase :		Optional[Any]								=    ['weaewolf']
			# S = 1, D = 0, I = 0, N = 8, CER = 0.125
			_UpperCamelCase :		List[str]								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.1_2_5) < 1e-6)


		def 		A__       (					self):
			_UpperCamelCase :		int								=    ['werewolf']
			_UpperCamelCase :		Any								=    ['wereawolf']
			# S = 0, D = 1, I = 0, N = 8, CER = 0.125
			_UpperCamelCase :		Optional[int]								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.1_2_5) < 1e-6)


		def 		A__       (					self):
			_UpperCamelCase :		str								=    ['werewolf']
			_UpperCamelCase :		Optional[Any]								=    ['wereolf']
			# S = 0, D = 0, I = 1, N = 8, CER = 0.125
			_UpperCamelCase :		List[Any]								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.1_2_5) < 1e-6)


		def 		A__       (					self):
			_UpperCamelCase :		Optional[Any]								=    ['werewolf']
			_UpperCamelCase :		Optional[Any]								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertEqual(__snake_case						,							0.0)


		def 		A__       (					self):
			_UpperCamelCase :		Union[str, Any]								=    ['werewolf', 'I am your father']
			_UpperCamelCase :		Dict								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertEqual(__snake_case						,							0.0)

			_UpperCamelCase :		List[Any]								=    ['werewolf', 'I am your father', 'doge']
			_UpperCamelCase :		Optional[int]								=    ['werxwolf', 'I       am your father', 'doge']
			# S = 1, D = 0, I = 0, N = 28, CER = 1 / 28
			_UpperCamelCase :		Any								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.0_3_5_7_1_4_2_8) < 1e-6)


		def 		A__       (					self):
			_UpperCamelCase :		Tuple								=    ['My hovercraft', 'is full of eels']
			_UpperCamelCase :		List[Any]								=    ['My hovercraft is full', ' of eels']
			# S = 0, D = 0, I = 2, N = 28, CER = 2 / 28
			# whitespace at the front of " of eels" will be strip during preporcessing
			# so need to insert 2 whitespaces
			_UpperCamelCase :		Dict								=    cer.compute(predictions=__snake_case						,							references=__snake_case						,							concatenate_texts=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.0_7_1_4_2_8) < 1e-6)


		def 		A__       (					self):
			_UpperCamelCase :		Union[str, Any]								=    ['我能吞下玻璃而不伤身体']
			_UpperCamelCase :		Optional[Any]								=    [' 能吞虾玻璃而 不霜身体啦']
			# S = 3, D = 2, I = 0, N = 11, CER = 5 / 11
			_UpperCamelCase :		Any								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.4_5_4_5_4_5_4_5_4_5) < 1e-6)

			_UpperCamelCase :		Union[str, Any]								=    ['我能吞下玻璃', '而不伤身体']
			_UpperCamelCase :		List[str]								=    ['我    能 吞 下 玻 璃', '而不伤身体']
			# S = 0, D = 5, I = 0, N = 11, CER = 5 / 11
			_UpperCamelCase :		Union[str, Any]								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertTrue(abs(char_error_rate - 0.4_5_4_5_4_5_4_5_4_5_4_5) < 1e-6)

			_UpperCamelCase :		Optional[int]								=    ['我能吞下玻璃而不伤身体']
			_UpperCamelCase :		int								=    cer.compute(predictions=__snake_case						,							references=__snake_case)
			self.assertFalse(__snake_case						,							0.0)






		def 		A__       (					self):
			_UpperCamelCase :		Tuple								=    ['']
			_UpperCamelCase :		Dict								=    ['Hypothesis']
			with self.assertRaises(__snake_case):
				cer.compute(predictions=__snake_case						,							references=__snake_case)


if __name__ == "__main__":
							unittest.main()




