

'''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 				A_					(	snake_case ,    snake_case=0 ,    snake_case=None	):
  return [" ".join(x.split(" "	)[begin:end]	) for x in answers if len(x.split(" "	)	) > 1]





def 				A_					(	snake_case ,    snake_case=False	):
  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:Any   =   (
        given_answers + get_sub_answers(snake_case ,    begin=1	) + get_sub_answers(snake_case ,    end=-1	)
    )
  SCREAMING_SNAKE_CASE:Optional[int]   =   []
  for answer in given_answers:
    SCREAMING_SNAKE_CASE:str   =   answer.replace("_" ,    " "	).lower()
    SCREAMING_SNAKE_CASE:Optional[int]   =   "".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 				A_					(	snake_case ,    snake_case ,    snake_case=1 ,    snake_case=100	):
  SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE:int   =   jax.lax.top_k(snake_case ,    snake_case	)
  SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE:List[Any]   =   jax.lax.top_k(snake_case ,    snake_case	)

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

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





def 				A_					(	snake_case	):
  SCREAMING_SNAKE_CASE:int   =   sample["question"]["text"]
  SCREAMING_SNAKE_CASE:List[Any]   =   sample["document"]["tokens"]["token"]
  SCREAMING_SNAKE_CASE:Optional[Any]   =   sample["document"]["tokens"]["is_html"]
  SCREAMING_SNAKE_CASE:Optional[Any]   =   sample["annotations"]["long_answer"]
  SCREAMING_SNAKE_CASE:str   =   sample["annotations"]["short_answers"]

  SCREAMING_SNAKE_CASE:Dict   =   " ".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:Dict   =   []
  for s in short_answers:
    short_targets.extend(s["text"]	)
  SCREAMING_SNAKE_CASE:List[str]   =   list(set(snake_case	)	)

  SCREAMING_SNAKE_CASE:Any   =   []
  for s in long_answers:
    if s["start_token"] == -1:
      continue
    SCREAMING_SNAKE_CASE:Tuple   =   context[s["start_token"] : s["end_token"]]
    SCREAMING_SNAKE_CASE:Tuple   =   is_html[s["start_token"] : s["end_token"]]
    SCREAMING_SNAKE_CASE:Optional[int]   =   " ".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:Union[str, Any]   =   "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_					(	):
  SCREAMING_SNAKE_CASE:Dict   =   load_from_disk("natural-questions-validation"	)
  SCREAMING_SNAKE_CASE:Optional[Any]   =   dataset.map(snake_case	).remove_columns(["annotations", "document", "id"]	)
  print(snake_case	)

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

  SCREAMING_SNAKE_CASE:str   =   "vasudevgupta/flax-bigbird-natural-questions"
  SCREAMING_SNAKE_CASE:List[str]   =   FlaxBigBirdForNaturalQuestions.from_pretrained(snake_case	)
  SCREAMING_SNAKE_CASE:Dict   =   BigBirdTokenizerFast.from_pretrained(snake_case	)

  @jax.jit
  def forward(*snake_case ,    **snake_case	):
    SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE:Union[str, Any]   =   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:Optional[Any]   =   tokenizer(
        example["question"] ,    example["context"] ,    return_tensors="np" ,    max_length=4096 ,    padding="max_length" ,    truncation=snake_case ,    )

    SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE:List[Any]   =   forward(**snake_case	)

    SCREAMING_SNAKE_CASE:str   =   CATEGORY_MAPPING[category.item()]

    SCREAMING_SNAKE_CASE:Optional[int]   =   example["long"] + example["short"]
    if example["category"] in ["yes", "no", "null"]:
      SCREAMING_SNAKE_CASE:Any   =   [example["category"]]
    SCREAMING_SNAKE_CASE:Optional[int]   =   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:List[Any]   =   [predicted_category]
      SCREAMING_SNAKE_CASE:Tuple   =   example["output"] == example["targets"]
      SCREAMING_SNAKE_CASE:int   =   predicted_category != "null"
      return example

    SCREAMING_SNAKE_CASE:List[Any]   =   38 if predicted_category == "short" else 1024
    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:List[str]   =   [tokenizer.decode(input_ids[start_score : end_score + 1]	)]

    SCREAMING_SNAKE_CASE:Optional[int]   =   expand_to_aliases(example["targets"] ,    make_sub_answers=snake_case	)
    SCREAMING_SNAKE_CASE:Union[str, Any]   =   expand_to_aliases(example["output"]	)

    # some preprocessing to both prediction and answer
    SCREAMING_SNAKE_CASE:int   =   {"".join(a.split()	) for a in answers}
    SCREAMING_SNAKE_CASE:List[str]   =   {"".join(p.split()	) for p in predictions}
    SCREAMING_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
    SCREAMING_SNAKE_CASE:Any   =   len(list(answers & predictions	)	) > 0
    SCREAMING_SNAKE_CASE:str   =   predicted_category != "null" and len(snake_case	) > 0

    return example

  SCREAMING_SNAKE_CASE:List[str]   =   short_validation_dataset.map(snake_case	)

  SCREAMING_SNAKE_CASE:Tuple   =   len(snake_case	)
  SCREAMING_SNAKE_CASE:Any   =   len(short_validation_dataset.filter(lambda snake_case	: x["match"] == 1	)	)
  print("EM score:" ,    (matched / total) * 100 ,    "%"	)


if __name__ == "__main__":
    main()









'''simple docstring'''

import shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
    from transformers import WhisperFeatureExtractor, WhisperProcessor


A_      =							5_03_58
A_      =							5_03_62





@require_torch
@require_torchaudio
@require_sentencepiece
class 							_snake_case	(    unittest.TestCase  ):
   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:int   =   "openai/whisper-small.en"
     SCREAMING_SNAKE_CASE:Optional[Any]   =   tempfile.mkdtemp()
   def       __UpperCamelCase  (	self     : Any       ,**SCREAMING_SNAKE_CASE__     : List[str]							):
     return WhisperTokenizer.from_pretrained(self.checkpoint       ,**SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : List[str]       ,**SCREAMING_SNAKE_CASE__     : Dict							):
     return WhisperFeatureExtractor.from_pretrained(self.checkpoint       ,**SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : str							):
     shutil.rmtree(self.tmpdirname							)
   def       __UpperCamelCase  (	self     : Any							):
     SCREAMING_SNAKE_CASE:Dict   =   self.get_tokenizer()
     SCREAMING_SNAKE_CASE:Optional[int]   =   self.get_feature_extractor()

     SCREAMING_SNAKE_CASE:Dict   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     processor.save_pretrained(self.tmpdirname							)
     SCREAMING_SNAKE_CASE:List[Any]   =   WhisperProcessor.from_pretrained(self.tmpdirname							)

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

     self.assertEqual(processor.feature_extractor.to_json_string()       ,feature_extractor.to_json_string()							)
     self.assertIsInstance(processor.feature_extractor       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:int   =   WhisperProcessor(tokenizer=self.get_tokenizer()       ,feature_extractor=self.get_feature_extractor()							)
     processor.save_pretrained(self.tmpdirname							)

     SCREAMING_SNAKE_CASE:List[str]   =   self.get_tokenizer(bos_token="(BOS)"       ,eos_token="(EOS)"							)
     SCREAMING_SNAKE_CASE:Tuple   =   self.get_feature_extractor(do_normalize=SCREAMING_SNAKE_CASE__       ,padding_value=1.0							)

     SCREAMING_SNAKE_CASE:Any   =   WhisperProcessor.from_pretrained(
         self.tmpdirname       ,bos_token="(BOS)"       ,eos_token="(EOS)"       ,do_normalize=SCREAMING_SNAKE_CASE__       ,padding_value=1.0							)

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

     self.assertEqual(processor.feature_extractor.to_json_string()       ,feature_extractor_add_kwargs.to_json_string()							)
     self.assertIsInstance(processor.feature_extractor       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : List[Any]							):
     SCREAMING_SNAKE_CASE:Any   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:List[Any]   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Optional[Any]   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Any   =   floats_list((3, 1_000)							)

     SCREAMING_SNAKE_CASE:Any   =   feature_extractor(SCREAMING_SNAKE_CASE__       ,return_tensors="np"							)
     SCREAMING_SNAKE_CASE:List[str]   =   processor(SCREAMING_SNAKE_CASE__       ,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       __UpperCamelCase  (	self     : List[Any]							):
     SCREAMING_SNAKE_CASE:int   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:Tuple   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:str   =   "This is a test string"

     SCREAMING_SNAKE_CASE:int   =   processor(text=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:str   =   tokenizer(SCREAMING_SNAKE_CASE__							)

     for key in encoded_tok.keys():
       self.assertListEqual(encoded_tok[key]       ,encoded_processor[key]							)
   def       __UpperCamelCase  (	self     : str							):
     SCREAMING_SNAKE_CASE:List[Any]   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:List[str]   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Optional[Any]   =   [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

     SCREAMING_SNAKE_CASE:Any   =   processor.batch_decode(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, Any]   =   tokenizer.batch_decode(SCREAMING_SNAKE_CASE__							)

     self.assertListEqual(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : int							):
     SCREAMING_SNAKE_CASE:Tuple   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:Optional[int]   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Union[str, Any]   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)

     self.assertListEqual(
         processor.model_input_names       ,feature_extractor.model_input_names       ,msg="`processor` and `feature_extractor` model input names do not match"       ,)
   def       __UpperCamelCase  (	self     : int							):
     SCREAMING_SNAKE_CASE:Tuple   =   self.get_feature_extractor()
     SCREAMING_SNAKE_CASE:Any   =   self.get_tokenizer()

     SCREAMING_SNAKE_CASE:Any   =   WhisperProcessor(tokenizer=SCREAMING_SNAKE_CASE__       ,feature_extractor=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Optional[int]   =   processor.get_decoder_prompt_ids(task="transcribe"       ,no_timestamps=SCREAMING_SNAKE_CASE__							)

     self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__							)
     for ids in forced_decoder_ids:
       self.assertIsInstance(SCREAMING_SNAKE_CASE__       ,(list, tuple)							)

     SCREAMING_SNAKE_CASE:Any   =   [TRANSCRIBE, NOTIMESTAMPS]
     self.assertListEqual([ids[-1] for ids in forced_decoder_ids]       ,SCREAMING_SNAKE_CASE__							)
   def       __UpperCamelCase  (	self     : Dict							):
     SCREAMING_SNAKE_CASE:Tuple   =   WhisperProcessor(tokenizer=self.get_tokenizer()       ,feature_extractor=self.get_feature_extractor()							)
     SCREAMING_SNAKE_CASE:Dict   =   processor.get_prompt_ids("Mr. Quilter"							)
     SCREAMING_SNAKE_CASE:int   =   processor.tokenizer.decode(SCREAMING_SNAKE_CASE__							)

     self.assertListEqual(prompt_ids.tolist()       ,[50_360, 1_770, 13, 2_264, 346, 353]							)
     self.assertEqual(SCREAMING_SNAKE_CASE__       ,"<|startofprev|> Mr. Quilter"							)
   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:Dict   =   WhisperProcessor(tokenizer=self.get_tokenizer()       ,feature_extractor=self.get_feature_extractor()							)
     SCREAMING_SNAKE_CASE:List[Any]   =   processor.get_prompt_ids(""							)
     SCREAMING_SNAKE_CASE:Tuple   =   processor.tokenizer.decode(SCREAMING_SNAKE_CASE__							)

     self.assertListEqual(prompt_ids.tolist()       ,[50_360, 220]							)
     self.assertEqual(SCREAMING_SNAKE_CASE__       ,"<|startofprev|> "							)







   def       __UpperCamelCase  (	self     : Union[str, Any]							):
     SCREAMING_SNAKE_CASE:Tuple   =   WhisperProcessor(tokenizer=self.get_tokenizer()       ,feature_extractor=self.get_feature_extractor()							)

     def _test_prompt_error_raised_helper(SCREAMING_SNAKE_CASE__     : int       ,SCREAMING_SNAKE_CASE__     : List[str]							):
       with pytest.raises(SCREAMING_SNAKE_CASE__							) as excinfo:
         processor.get_prompt_ids(SCREAMING_SNAKE_CASE__							)
       SCREAMING_SNAKE_CASE:List[Any]   =   F'''Encountered text in the prompt corresponding to disallowed special token: {special_token}.'''
       self.assertEqual(SCREAMING_SNAKE_CASE__       ,str(excinfo.value							)							)

     _test_prompt_error_raised_helper("<|startofprev|> test"       ,"<|startofprev|>"							)
     _test_prompt_error_raised_helper("test <|notimestamps|>"       ,"<|notimestamps|>"							)
     _test_prompt_error_raised_helper("test <|zh|> test <|transcribe|>"       ,"<|zh|>"							)






