






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

from transformers import BigBirdTokenizerFast


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


def   _A       (    SCREAMING_SNAKE_CASE__  :     Optional[int]					,  SCREAMING_SNAKE_CASE__  :     Union[str, Any]=0					,  SCREAMING_SNAKE_CASE__  :     List[str]=None	):
      return [" ".join(x.split(''' '''	)[begin:end]	) for x in answers if len(x.split(''' '''	)	) > 1]


def   _A       (    SCREAMING_SNAKE_CASE__  :     List[Any]					,  SCREAMING_SNAKE_CASE__  :     Tuple=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"
            UpperCamelCase	:Union[str, Any]     							=       (
                given_answers + get_sub_answers(SCREAMING_SNAKE_CASE__					,  begin=1	) + get_sub_answers(SCREAMING_SNAKE_CASE__					,  end=-1	)
            )
      UpperCamelCase	:Union[str, Any]     							=       []
      for answer in given_answers:
            UpperCamelCase	:List[Any]     							=       answer.replace('''_'''					,  ''' '''	).lower()
            UpperCamelCase	: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   _A       (    SCREAMING_SNAKE_CASE__  :     List[str]					,  SCREAMING_SNAKE_CASE__  :     int					,  SCREAMING_SNAKE_CASE__  :     Dict=1					,  SCREAMING_SNAKE_CASE__  :     Tuple=100	):
      UpperCamelCase	:Optional[Any]     							=       jax.lax.top_k(SCREAMING_SNAKE_CASE__					,  SCREAMING_SNAKE_CASE__	)
      UpperCamelCase	:List[Any]     							=       jax.lax.top_k(SCREAMING_SNAKE_CASE__					,  SCREAMING_SNAKE_CASE__	)

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

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


def   _A       (    SCREAMING_SNAKE_CASE__  :     List[str]	):
      UpperCamelCase	:List[str]     							=       sample['''question''']['''text''']
      UpperCamelCase	:str     							=       sample['''document''']['''tokens''']['''token''']
      UpperCamelCase	:Union[str, Any]     							=       sample['''document''']['''tokens''']['''is_html''']
      UpperCamelCase	:Any     							=       sample['''annotations''']['''long_answer''']
      UpperCamelCase	:Optional[int]     							=       sample['''annotations''']['''short_answers''']

      UpperCamelCase	:List[Any]     							=       ''' '''.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",
                  }

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

      UpperCamelCase	:List[str]     							=       []
      for s in long_answers:
            if s["start_token"] == -1:
                  continue
            UpperCamelCase	:List[str]     							=       context[s['''start_token'''] : s['''end_token''']]
            UpperCamelCase	:Dict     							=       is_html[s['''start_token'''] : s['''end_token''']]
            UpperCamelCase	:Union[str, 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__	)

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

      UpperCamelCase	:List[str]     							=       dataset.filter(lambda SCREAMING_SNAKE_CASE__	: (len(x['''question''']	) + len(x['''context''']	)) < 4 * 4096	)
      UpperCamelCase	:Dict     							=       short_validation_dataset.filter(lambda SCREAMING_SNAKE_CASE__	: x["category"] != "null"	)
      short_validation_dataset

      UpperCamelCase	:Union[str, Any]     							=       '''vasudevgupta/flax-bigbird-natural-questions'''
      UpperCamelCase	:Union[str, Any]     							=       FlaxBigBirdForNaturalQuestions.from_pretrained(SCREAMING_SNAKE_CASE__	)
      UpperCamelCase	:Dict     							=       BigBirdTokenizerFast.from_pretrained(SCREAMING_SNAKE_CASE__	)

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

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

            UpperCamelCase	:Any     							=       forward(**SCREAMING_SNAKE_CASE__	)

            UpperCamelCase	:Tuple     							=       CATEGORY_MAPPING[category.item()]

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

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

            UpperCamelCase	:Dict     							=       38 if predicted_category == '''short''' else 1024
            UpperCamelCase	:Optional[Any]     							=       get_best_valid_start_end_idx(
                start_scores[0]					,  end_scores[0]					,  top_k=8					,  max_size=SCREAMING_SNAKE_CASE__	)

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

            UpperCamelCase	:Optional[int]     							=       expand_to_aliases(example['''targets''']					,  make_sub_answers=SCREAMING_SNAKE_CASE__	)
            UpperCamelCase	:int     							=       expand_to_aliases(example['''output''']	)

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

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

            return example

      UpperCamelCase	:Union[str, Any]     							=       short_validation_dataset.map(SCREAMING_SNAKE_CASE__	)

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


if __name__ == "__main__":
 main()













import warnings
from collections.abc import Callable
from typing import Any

import numpy
from matplotlib import pyplot

__snake_case							 =  0.2_5 + 0.0j
__snake_case							 =  -0.4 + 0.6j
__snake_case							 =  -0.1 + 0.651j
__snake_case							 =  -2.0
__snake_case							 =  56
__snake_case							 =  2.0
__snake_case							 =  6_66


def   _A       (    SCREAMING_SNAKE_CASE__  :     complex					,  SCREAMING_SNAKE_CASE__  :     numpy.ndarray	):
    return numpy.exp(SCREAMING_SNAKE_CASE__	) + c_parameter


def   _A       (    SCREAMING_SNAKE_CASE__  :     complex					,  SCREAMING_SNAKE_CASE__  :     numpy.ndarray	):
    return z_values * z_values + c_parameter


def   _A       (    SCREAMING_SNAKE_CASE__  :     float					,  SCREAMING_SNAKE_CASE__  :     int	):
    UpperCamelCase	:List[Any]     							=       numpy.linspace(-window_size					,  SCREAMING_SNAKE_CASE__					,  SCREAMING_SNAKE_CASE__	)
    UpperCamelCase	:Any     							=       x.reshape((nb_pixels, 1)	)
    UpperCamelCase	:Tuple     							=       numpy.linspace(-window_size					,  SCREAMING_SNAKE_CASE__					,  SCREAMING_SNAKE_CASE__	)
    UpperCamelCase	:int     							=       y.reshape((1, nb_pixels)	)
    return x + 1.0J * y


def   _A       (    SCREAMING_SNAKE_CASE__  :     Callable[[Any, numpy.ndarray], numpy.ndarray]					,  SCREAMING_SNAKE_CASE__  :     Any					,  SCREAMING_SNAKE_CASE__  :     int					,  SCREAMING_SNAKE_CASE__  :     numpy.ndarray					,  SCREAMING_SNAKE_CASE__  :     float | None = None					,  ):
    UpperCamelCase	:Dict     							=       z_a.astype('''complex64'''	)
    for _ in range(SCREAMING_SNAKE_CASE__	):
        UpperCamelCase	:int     							=       eval_function(SCREAMING_SNAKE_CASE__					,  SCREAMING_SNAKE_CASE__	)
        if infinity is not None:
            numpy.nan_to_num(SCREAMING_SNAKE_CASE__					,  copy=SCREAMING_SNAKE_CASE__					,  nan=SCREAMING_SNAKE_CASE__	)
            UpperCamelCase	:Tuple     							=       infinity
    return z_n


def   _A       (    SCREAMING_SNAKE_CASE__  :     str					,  SCREAMING_SNAKE_CASE__  :     Any					,  SCREAMING_SNAKE_CASE__  :     float					,  SCREAMING_SNAKE_CASE__  :     numpy.ndarray					,  ):
    UpperCamelCase	:Any     							=       (abs(SCREAMING_SNAKE_CASE__	)).transpose()
    UpperCamelCase	:Any     							=       abs_z_final[::-1, :]
    pyplot.matshow(abs_z_final < escape_radius	)
    pyplot.title(F'''Julia set of ${function_label}$, $c={function_params}$'''	)
    pyplot.show()


def   _A       (    ):
    warnings.filterwarnings(
        '''ignore'''					,  category=SCREAMING_SNAKE_CASE__					,  message='''overflow encountered in multiply'''	)
    warnings.filterwarnings(
        '''ignore'''					,  category=SCREAMING_SNAKE_CASE__					,  message='''invalid value encountered in multiply'''					,  )
    warnings.filterwarnings(
        '''ignore'''					,  category=SCREAMING_SNAKE_CASE__					,  message='''overflow encountered in absolute'''	)
    warnings.filterwarnings(
        '''ignore'''					,  category=SCREAMING_SNAKE_CASE__					,  message='''overflow encountered in exp'''	)


if __name__ == "__main__":
 __snake_case							 =  prepare_grid(window_size, nb_pixels)

 ignore_overflow_warnings()  # See file header for explanations

 __snake_case							 =  24
 __snake_case							 =  2 * abs(c_cauliflower) + 1
 __snake_case							 =  iterate_function(
     eval_quadratic_polynomial,
     c_cauliflower,
     nb_iterations,
     z_a,
     infinity=1.1 * escape_radius,
 )
 show_results("""z^2+c""", c_cauliflower, escape_radius, z_final)

 __snake_case							 =  64
 __snake_case							 =  2 * abs(c_polynomial_a) + 1
 __snake_case							 =  iterate_function(
     eval_quadratic_polynomial,
     c_polynomial_a,
     nb_iterations,
     z_a,
     infinity=1.1 * escape_radius,
 )
 show_results("""z^2+c""", c_polynomial_a, escape_radius, z_final)

 __snake_case							 =  1_61
 __snake_case							 =  2 * abs(c_polynomial_a) + 1
 __snake_case							 =  iterate_function(
     eval_quadratic_polynomial,
     c_polynomial_a,
     nb_iterations,
     z_a,
     infinity=1.1 * escape_radius,
 )
 show_results("""z^2+c""", c_polynomial_a, escape_radius, z_final)

 __snake_case							 =  12
 __snake_case							 =  1_0_0_0_0.0
 __snake_case							 =  iterate_function(
     eval_exponential,
     c_exponential,
     nb_iterations,
     z_a + 2,
     infinity=1.0E1_0,
 )
 show_results("""e^z+c""", c_exponential, escape_radius, z_final)





