




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






def 		UpperCamelCase__			(		__SCREAMING_SNAKE_CASE							,			__SCREAMING_SNAKE_CASE=0							,			__SCREAMING_SNAKE_CASE=None		)						->      Union[str, Any]:
      return [" ".join(x.split(' '		)[begin:end]		) for x in answers if len(x.split(' '		)		) > 1]






def 		UpperCamelCase__			(		__SCREAMING_SNAKE_CASE							,			__SCREAMING_SNAKE_CASE=False		)						->      Union[str, Any]:
      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"
            snake_case__					:      Optional[Any]  			=						(
                given_answers + get_sub_answers(__SCREAMING_SNAKE_CASE							,			begin=1		) + get_sub_answers(__SCREAMING_SNAKE_CASE							,			end=-1		)
            )
      snake_case__					:      List[str]  			=						[]
      for answer in given_answers:
            snake_case__					:      List[str]  			=						answer.replace('_'							,			' '		).lower()
            snake_case__					:      List[str]  			=						''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ' ' for c in alias		)
            answers.append(' '.join(alias.split()		).strip()		)
      return set(__SCREAMING_SNAKE_CASE		)






def 		UpperCamelCase__			(		__SCREAMING_SNAKE_CASE							,			__SCREAMING_SNAKE_CASE							,			__SCREAMING_SNAKE_CASE=1							,			__SCREAMING_SNAKE_CASE=100		)						->      Optional[Any]:
      snake_case__			,       snake_case__					:      List[Any]  			=						jax.lax.top_k(__SCREAMING_SNAKE_CASE							,			__SCREAMING_SNAKE_CASE		)
      snake_case__			,       snake_case__					:      List[Any]  			=						jax.lax.top_k(__SCREAMING_SNAKE_CASE							,			__SCREAMING_SNAKE_CASE		)

      snake_case__					:      Optional[Any]  			=						best_end_idx[:, None] - best_start_idx[None, :]
      snake_case__					:      List[Any]  			=						jnp.logical_or(widths < 0							,			widths > max_size		)
      snake_case__					:      Union[str, Any]  			=						(best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
      snake_case__					:      List[Any]  			=						jnp.argmax(__SCREAMING_SNAKE_CASE		).item()

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






def 		UpperCamelCase__			(		__SCREAMING_SNAKE_CASE		)						->      List[str]:
      snake_case__					:      Optional[Any]  			=						sample['question']['text']
      snake_case__					:      List[Any]  			=						sample['document']['tokens']['token']
      snake_case__					:      List[str]  			=						sample['document']['tokens']['is_html']
      snake_case__					:      Optional[Any]  			=						sample['annotations']['long_answer']
      snake_case__					:      List[Any]  			=						sample['annotations']['short_answers']

      snake_case__					:      Dict  			=						' '.join([context[i] for i in range(len(__SCREAMING_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",
                  }

      snake_case__					:      List[Any]  			=						[]
      for s in short_answers:
            short_targets.extend(s['text']		)
      snake_case__					:      List[str]  			=						list(set(__SCREAMING_SNAKE_CASE		)		)

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

      snake_case__					:      Tuple  			=						'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__			(		)						->      Optional[int]:
      snake_case__					:      Optional[Any]  			=						load_from_disk('natural-questions-validation'		)
      snake_case__					:      Optional[int]  			=						dataset.map(__SCREAMING_SNAKE_CASE		).remove_columns(['annotations', 'document', 'id']		)
      print(__SCREAMING_SNAKE_CASE		)

      snake_case__					:      Tuple  			=						dataset.filter(lambda __SCREAMING_SNAKE_CASE		: (len(x['question']		) + len(x['context']		)) < 4 * 4096		)
      snake_case__					:      List[Any]  			=						short_validation_dataset.filter(lambda __SCREAMING_SNAKE_CASE		: x["category"] != "null"		)
      short_validation_dataset

      snake_case__					:      Any  			=						'vasudevgupta/flax-bigbird-natural-questions'
      snake_case__					:      Optional[Any]  			=						FlaxBigBirdForNaturalQuestions.from_pretrained(__SCREAMING_SNAKE_CASE		)
      snake_case__					:      Tuple  			=						BigBirdTokenizerFast.from_pretrained(__SCREAMING_SNAKE_CASE		)

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

      def evaluate(__SCREAMING_SNAKE_CASE		):
            # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
            snake_case__					:      Union[str, Any]  			=						tokenizer(
                example['question']							,			example['context']							,			return_tensors='np'							,			max_length=4096							,			padding='max_length'							,			truncation=__SCREAMING_SNAKE_CASE							,			)

            snake_case__			,       snake_case__			,       snake_case__					:      Optional[int]  			=						forward(**__SCREAMING_SNAKE_CASE		)

            snake_case__					:      Any  			=						CATEGORY_MAPPING[category.item()]

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

            if predicted_category in ["yes", "no", "null"]:
                  snake_case__					:      Tuple  			=						[predicted_category]
                  snake_case__					:      Optional[int]  			=						example['output'] == example['targets']
                  snake_case__					:      Optional[int]  			=						predicted_category != 'null'
                  return example

            snake_case__					:      Tuple  			=						38 if predicted_category == 'short' else 1024
            snake_case__			,       snake_case__					:      int  			=						get_best_valid_start_end_idx(
                start_scores[0]							,			end_scores[0]							,			top_k=8							,			max_size=__SCREAMING_SNAKE_CASE		)

            snake_case__					:      List[str]  			=						inputs['input_ids'][0].tolist()
            snake_case__					:      List[Any]  			=						[tokenizer.decode(input_ids[start_score : end_score + 1]		)]

            snake_case__					:      Dict  			=						expand_to_aliases(example['targets']							,			make_sub_answers=__SCREAMING_SNAKE_CASE		)
            snake_case__					:      Optional[int]  			=						expand_to_aliases(example['output']		)

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

            # if there is a common element, it's a exact match
            snake_case__					:      str  			=						len(list(answers & predictions		)		) > 0
            snake_case__					:      List[str]  			=						predicted_category != 'null' and len(__SCREAMING_SNAKE_CASE		) > 0

            return example

      snake_case__					:      List[Any]  			=						short_validation_dataset.map(__SCREAMING_SNAKE_CASE		)

      snake_case__					:      Union[str, Any]  			=						len(__SCREAMING_SNAKE_CASE		)
      snake_case__					:      Tuple  			=						len(short_validation_dataset.filter(lambda __SCREAMING_SNAKE_CASE		: x["match"] == 1		)		)
      print('EM score:'							,			(matched / total) * 100							,			'%'		)


if __name__ == "__main__":
      main()






'''simple docstring'''






import shutil
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SpeechaTextTokenizer, is_speech_available
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_torch, require_torchaudio
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_speech_to_text import floats_list


if is_speech_available():
      from transformers import SpeechaTextFeatureExtractor, SpeechaTextProcessor


A_ 			=		get_tests_dir("fixtures/test_sentencepiece.model")






@require_torch
@require_torchaudio
@require_sentencepiece
class 					lowercase_			(	unittest.TestCase       ):





     def  _lowerCAmelCase			(    self    :				List[str] ):
           snake_case__					:      Dict  			=						tempfile.mkdtemp()

           snake_case__					:      Union[str, Any]  			=						['<s>', '<pad>', '</s>', '<unk>', '▁This', '▁is', '▁a', '▁t', 'est']
           snake_case__					:      str  			=						dict(zip(__lowerCamelCase    ,       range(len(__lowerCamelCase ) ) ) )
           snake_case__					:      Union[str, Any]  			=						Path(self.tmpdirname )
           save_json(__lowerCamelCase    ,       save_dir / VOCAB_FILES_NAMES['vocab_file'] )
           if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
                 copyfile(__lowerCamelCase    ,       save_dir / VOCAB_FILES_NAMES['spm_file'] )

           snake_case__					:      Optional[Any]  			=						SpeechaTextTokenizer.from_pretrained(self.tmpdirname )
           tokenizer.save_pretrained(self.tmpdirname )

           snake_case__					:      str  			=						{
               'feature_size': 24,
               'num_mel_bins': 24,
               'padding_value': 0.0,
               'sampling_rate': 16000,
               'return_attention_mask': False,
               'do_normalize': True,
           }
           save_json(__lowerCamelCase    ,       save_dir / FEATURE_EXTRACTOR_NAME )





     def  _lowerCAmelCase			(    self    :				str    ,       **__lowerCamelCase    :				int ):
           return SpeechaTextTokenizer.from_pretrained(self.tmpdirname    ,       **__lowerCamelCase )





     def  _lowerCAmelCase			(    self    :				int    ,       **__lowerCamelCase    :				List[Any] ):
           return SpeechaTextFeatureExtractor.from_pretrained(self.tmpdirname    ,       **__lowerCamelCase )





     def  _lowerCAmelCase			(    self    :				List[Any] ):
           shutil.rmtree(self.tmpdirname )





     def  _lowerCAmelCase			(    self    :				Dict ):
           snake_case__					:      Dict  			=						self.get_tokenizer()
           snake_case__					:      Any  			=						self.get_feature_extractor()

           snake_case__					:      Tuple  			=						SpeechaTextProcessor(tokenizer=__lowerCamelCase    ,       feature_extractor=__lowerCamelCase )

           processor.save_pretrained(self.tmpdirname )
           snake_case__					:      Any  			=						SpeechaTextProcessor.from_pretrained(self.tmpdirname )

           self.assertEqual(processor.tokenizer.get_vocab()    ,       tokenizer.get_vocab() )
           self.assertIsInstance(processor.tokenizer    ,       __lowerCamelCase )

           self.assertEqual(processor.feature_extractor.to_json_string()    ,       feature_extractor.to_json_string() )
           self.assertIsInstance(processor.feature_extractor    ,       __lowerCamelCase )





     def  _lowerCAmelCase			(    self    :				List[Any] ):
           snake_case__					:      Any  			=						SpeechaTextProcessor(
               tokenizer=self.get_tokenizer()    ,       feature_extractor=self.get_feature_extractor() )
           processor.save_pretrained(self.tmpdirname )

           snake_case__					:      Tuple  			=						self.get_tokenizer(bos_token='(BOS)'    ,       eos_token='(EOS)' )
           snake_case__					:      int  			=						self.get_feature_extractor(do_normalize=__lowerCamelCase    ,       padding_value=1.0 )

           snake_case__					:      Any  			=						SpeechaTextProcessor.from_pretrained(
               self.tmpdirname    ,       bos_token='(BOS)'    ,       eos_token='(EOS)'    ,       do_normalize=__lowerCamelCase    ,       padding_value=1.0 )

           self.assertEqual(processor.tokenizer.get_vocab()    ,       tokenizer_add_kwargs.get_vocab() )
           self.assertIsInstance(processor.tokenizer    ,       __lowerCamelCase )

           self.assertEqual(processor.feature_extractor.to_json_string()    ,       feature_extractor_add_kwargs.to_json_string() )
           self.assertIsInstance(processor.feature_extractor    ,       __lowerCamelCase )





     def  _lowerCAmelCase			(    self    :				Any ):
           snake_case__					:      int  			=						self.get_feature_extractor()
           snake_case__					:      Optional[Any]  			=						self.get_tokenizer()

           snake_case__					:      List[str]  			=						SpeechaTextProcessor(tokenizer=__lowerCamelCase    ,       feature_extractor=__lowerCamelCase )

           snake_case__					:      List[str]  			=						floats_list((3, 1000) )

           snake_case__					:      Tuple  			=						feature_extractor(__lowerCamelCase    ,       return_tensors='np' )
           snake_case__					:      str  			=						processor(__lowerCamelCase    ,       return_tensors='np' )

           for key in input_feat_extract.keys():
                 self.assertAlmostEqual(input_feat_extract[key].sum()    ,       input_processor[key].sum()    ,       delta=1E-2 )





     def  _lowerCAmelCase			(    self    :				Dict ):
           snake_case__					:      Optional[int]  			=						self.get_feature_extractor()
           snake_case__					:      Optional[Any]  			=						self.get_tokenizer()

           snake_case__					:      Tuple  			=						SpeechaTextProcessor(tokenizer=__lowerCamelCase    ,       feature_extractor=__lowerCamelCase )

           snake_case__					:      Optional[int]  			=						'This is a test string'

           snake_case__					:      int  			=						processor(text=__lowerCamelCase )

           snake_case__					:      Tuple  			=						tokenizer(__lowerCamelCase )

           for key in encoded_tok.keys():
                 self.assertListEqual(encoded_tok[key]    ,       encoded_processor[key] )





     def  _lowerCAmelCase			(    self    :				int ):
           snake_case__					:      Optional[int]  			=						self.get_feature_extractor()
           snake_case__					:      int  			=						self.get_tokenizer()

           snake_case__					:      str  			=						SpeechaTextProcessor(tokenizer=__lowerCamelCase    ,       feature_extractor=__lowerCamelCase )

           snake_case__					:      str  			=						[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

           snake_case__					:      List[str]  			=						processor.batch_decode(__lowerCamelCase )
           snake_case__					:      Optional[int]  			=						tokenizer.batch_decode(__lowerCamelCase )

           self.assertListEqual(__lowerCamelCase    ,       __lowerCamelCase )





     def  _lowerCAmelCase			(    self    :				str ):
           snake_case__					:      List[str]  			=						self.get_feature_extractor()
           snake_case__					:      int  			=						self.get_tokenizer()

           snake_case__					:      int  			=						SpeechaTextProcessor(tokenizer=__lowerCamelCase    ,       feature_extractor=__lowerCamelCase )

           self.assertListEqual(
               processor.model_input_names    ,       feature_extractor.model_input_names    ,       msg='`processor` and `feature_extractor` model input names do not match'    ,       )
