




'''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__								=						{0: '''null''', 1: '''short''', 2: '''long''', 3: '''yes''', 4: '''no'''}
lowercase__								=						set(''''''.join(['''‘''', '''’''', '''´''', '''`''', '''.''', ''',''', '''-''', '''"''']))


def 					UpperCamelCase						(	snake_case						,  snake_case=0						,  snake_case=None ):




   '''simple docstring'''

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


def 					UpperCamelCase						(	snake_case						,  snake_case=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"
      SCREAMING_SNAKE_CASE__      :					Tuple							    =			(
          given_answers + get_sub_answers(snake_case						,  begin=1 ) + get_sub_answers(snake_case						,  end=-1 )
      )
   SCREAMING_SNAKE_CASE__      :					Any							    =			[]
   for answer in given_answers:
      SCREAMING_SNAKE_CASE__      :					int							    =			answer.replace('''_'''						,  ''' ''' ).lower()
      SCREAMING_SNAKE_CASE__      :					Tuple							    =			''''''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ''' ''' for c in alias )
      answers.append(''' '''.join(alias.split() ).strip() )
   return set(snake_case )


def 					UpperCamelCase						(	snake_case						,  snake_case						,  snake_case=1						,  snake_case=1_00 ):




   '''simple docstring'''

   SCREAMING_SNAKE_CASE__    ,					SCREAMING_SNAKE_CASE__      :					List[Any]							    =			jax.lax.top_k(snake_case						,  snake_case )
   SCREAMING_SNAKE_CASE__    ,					SCREAMING_SNAKE_CASE__      :					Optional[int]							    =			jax.lax.top_k(snake_case						,  snake_case )

   SCREAMING_SNAKE_CASE__      :					Optional[int]							    =			best_end_idx[:, None] - best_start_idx[None, :]
   SCREAMING_SNAKE_CASE__      :					Dict							    =			jnp.logical_or(widths < 0						,  widths > max_size )
   SCREAMING_SNAKE_CASE__      :					Any							    =			(best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
   SCREAMING_SNAKE_CASE__      :					str							    =			jnp.argmax(snake_case ).item()

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


def 					UpperCamelCase						(	snake_case ):




   '''simple docstring'''

   SCREAMING_SNAKE_CASE__      :					Any							    =			sample['''question''']['''text''']
   SCREAMING_SNAKE_CASE__      :					Optional[int]							    =			sample['''document''']['''tokens''']['''token''']
   SCREAMING_SNAKE_CASE__      :					Tuple							    =			sample['''document''']['''tokens''']['''is_html''']
   SCREAMING_SNAKE_CASE__      :					Tuple							    =			sample['''annotations''']['''long_answer''']
   SCREAMING_SNAKE_CASE__      :					str							    =			sample['''annotations''']['''short_answers''']

   SCREAMING_SNAKE_CASE__      :					List[str]							    =			''' '''.join([context[i] for i in range(len(snake_case ) ) 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",
         }

   SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			[]
   for s in short_answers:
      short_targets.extend(s['''text'''] )
   SCREAMING_SNAKE_CASE__      :					Dict							    =			list(set(snake_case ) )

   SCREAMING_SNAKE_CASE__      :					Any							    =			[]
   for s in long_answers:
      if s["start_token"] == -1:
         continue
      SCREAMING_SNAKE_CASE__      :					List[Any]							    =			context[s['''start_token'''] : s['''end_token''']]
      SCREAMING_SNAKE_CASE__      :					Dict							    =			is_html[s['''start_token'''] : s['''end_token''']]
      SCREAMING_SNAKE_CASE__      :					List[str]							    =			''' '''.join([answer[i] for i in range(len(snake_case ) ) if not html[i]] )
      if new_answer not in long_targets:
         long_targets.append(snake_case )

   SCREAMING_SNAKE_CASE__      :					int							    =			'''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						(	):




   '''simple docstring'''

   SCREAMING_SNAKE_CASE__      :					List[Any]							    =			load_from_disk('''natural-questions-validation''' )
   SCREAMING_SNAKE_CASE__      :					List[Any]							    =			dataset.map(snake_case ).remove_columns(['''annotations''', '''document''', '''id'''] )
   print(snake_case )

   SCREAMING_SNAKE_CASE__      :					Dict							    =			dataset.filter(lambda snake_case : (len(x['''question'''] ) + len(x['''context'''] )) < 4 * 40_96 )
   SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			short_validation_dataset.filter(lambda snake_case : x["category"] != "null" )
   short_validation_dataset

   SCREAMING_SNAKE_CASE__      :					List[Any]							    =			'''vasudevgupta/flax-bigbird-natural-questions'''
   SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			FlaxBigBirdForNaturalQuestions.from_pretrained(snake_case )
   SCREAMING_SNAKE_CASE__      :					List[str]							    =			BigBirdTokenizerFast.from_pretrained(snake_case )

   @jax.jit
   def forward(*snake_case						,  **snake_case ):
      SCREAMING_SNAKE_CASE__    ,					SCREAMING_SNAKE_CASE__    ,					SCREAMING_SNAKE_CASE__      :					str							    =			model(*snake_case						,  **snake_case )
      return start_logits, end_logits, jnp.argmax(snake_case						,  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
      SCREAMING_SNAKE_CASE__      :					Any							    =			tokenizer(
          example['''question''']						,  example['''context''']						,  return_tensors='''np'''						,  max_length=40_96						,  padding='''max_length'''						,  truncation=snake_case						,  )

      SCREAMING_SNAKE_CASE__    ,					SCREAMING_SNAKE_CASE__    ,					SCREAMING_SNAKE_CASE__      :					Any							    =			forward(**snake_case )

      SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			CATEGORY_MAPPING[category.item()]

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

      if predicted_category in ["yes", "no", "null"]:
         SCREAMING_SNAKE_CASE__      :					str							    =			[predicted_category]
         SCREAMING_SNAKE_CASE__      :					Tuple							    =			example['''output'''] == example['''targets''']
         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			predicted_category != '''null'''
         return example

      SCREAMING_SNAKE_CASE__      :					Any							    =			38 if predicted_category == '''short''' else 10_24
      SCREAMING_SNAKE_CASE__    ,					SCREAMING_SNAKE_CASE__      :					Tuple							    =			get_best_valid_start_end_idx(
          start_scores[0]						,  end_scores[0]						,  top_k=8						,  max_size=snake_case )

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

      SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			expand_to_aliases(example['''targets''']						,  make_sub_answers=snake_case )
      SCREAMING_SNAKE_CASE__      :					Any							    =			expand_to_aliases(example['''output'''] )

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

      # if there is a common element, it's a exact match
      SCREAMING_SNAKE_CASE__      :					Tuple							    =			len(list(answers & predictions ) ) > 0
      SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			predicted_category != '''null''' and len(snake_case ) > 0

      return example

   SCREAMING_SNAKE_CASE__      :					Optional[int]							    =			short_validation_dataset.map(snake_case )

   SCREAMING_SNAKE_CASE__      :					int							    =			len(snake_case )
   SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			len(short_validation_dataset.filter(lambda snake_case : x["match"] == 1 ) )
   print('''EM score:'''						,  (matched / total) * 1_00						,  '''%''' )


if __name__ == "__main__":
       main()







'''simple docstring'''





import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType






class     lowercase  ( UpperCAmelCase_):
      __lowercase					=	['''image_processor''', '''tokenizer''']
      __lowercase					=	'''LayoutLMv2ImageProcessor'''
      __lowercase					=	('''LayoutLMv2Tokenizer''', '''LayoutLMv2TokenizerFast''')






      def __init__(self			,	a__=None			,	a__=None			,	**a__       )							-> Any:
         SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			None
         if "feature_extractor" in kwargs:
            warnings.warn(
                '''The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'''
                ''' instead.'''			,	a__			,	)
            SCREAMING_SNAKE_CASE__      :					List[Any]							    =			kwargs.pop('''feature_extractor'''       )

         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			image_processor if image_processor is not None else feature_extractor
         if image_processor is None:
            raise ValueError('''You need to specify an `image_processor`.'''       )
         if tokenizer is None:
            raise ValueError('''You need to specify a `tokenizer`.'''       )

         super().__init__(a__			,	a__       )






      def __call__(self			,	a__			,	a__ = None			,	a__ = None			,	a__ = None			,	a__ = None			,	a__ = True			,	a__ = False			,	a__ = False			,	a__ = None			,	a__ = 0			,	a__ = None			,	a__ = None			,	a__ = None			,	a__ = False			,	a__ = False			,	a__ = False			,	a__ = False			,	a__ = True			,	a__ = None			,	**a__			,	)							-> BatchEncoding:
         # verify input
         if self.image_processor.apply_ocr and (boxes is not None):
            raise ValueError(
                '''You cannot provide bounding boxes if you initialized the image processor with apply_ocr set to True.'''       )

         if self.image_processor.apply_ocr and (word_labels is not None):
            raise ValueError(
                '''You cannot provide word labels if you initialized the image processor with apply_ocr set to True.'''       )

         if return_overflowing_tokens is True and return_offsets_mapping is False:
            raise ValueError('''You cannot return overflowing tokens without returning the offsets mapping.'''       )

         # first, apply the image processor
         SCREAMING_SNAKE_CASE__      :					str							    =			self.image_processor(images=a__			,	return_tensors=a__       )

         # second, apply the tokenizer
         if text is not None and self.image_processor.apply_ocr and text_pair is None:
            if isinstance(a__			,	a__       ):
               SCREAMING_SNAKE_CASE__      :					Tuple							    =			[text]  # add batch dimension (as the image processor always adds a batch dimension)
            SCREAMING_SNAKE_CASE__      :					Any							    =			features['''words''']

         SCREAMING_SNAKE_CASE__      :					int							    =			self.tokenizer(
             text=text if text is not None else features['''words''']			,	text_pair=text_pair if text_pair is not None else None			,	boxes=boxes if boxes is not None else features['''boxes''']			,	word_labels=a__			,	add_special_tokens=a__			,	padding=a__			,	truncation=a__			,	max_length=a__			,	stride=a__			,	pad_to_multiple_of=a__			,	return_token_type_ids=a__			,	return_attention_mask=a__			,	return_overflowing_tokens=a__			,	return_special_tokens_mask=a__			,	return_offsets_mapping=a__			,	return_length=a__			,	verbose=a__			,	return_tensors=a__			,	**a__			,	)

         # add pixel values
         SCREAMING_SNAKE_CASE__      :					Tuple							    =			features.pop('''pixel_values'''       )
         if return_overflowing_tokens is True:
            SCREAMING_SNAKE_CASE__      :					Tuple							    =			self.get_overflowing_images(a__			,	encoded_inputs['''overflow_to_sample_mapping''']       )
         SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			images

         return encoded_inputs






      def  A__(self			,	a__			,	a__       )							-> str:
         # in case there's an overflow, ensure each `input_ids` sample is mapped to its corresponding image
         SCREAMING_SNAKE_CASE__      :					int							    =			[]
         for sample_idx in overflow_to_sample_mapping:
            images_with_overflow.append(images[sample_idx]       )

         if len(a__       ) != len(a__       ):
            raise ValueError(
                '''Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got'''
                f""" {len(a__       )} and {len(a__       )}"""       )

         return images_with_overflow






      def  A__(self			,	*a__			,	**a__       )							-> List[Any]:
         return self.tokenizer.batch_decode(*a__			,	**a__       )






      def  A__(self			,	*a__			,	**a__       )							-> Optional[Any]:
         return self.tokenizer.decode(*a__			,	**a__       )






      @property
      def  A__(self       )							-> Tuple:
         return ["input_ids", "bbox", "token_type_ids", "attention_mask", "image"]






      @property
      def  A__(self       )							-> str:
         warnings.warn(
             '''`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'''			,	a__			,	)
         return self.image_processor_class






      @property
      def  A__(self       )							-> Dict:
         warnings.warn(
             '''`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'''			,	a__			,	)
         return self.image_processor

