




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

from transformers import BigBirdTokenizerFast


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



def     lowerCAmelCase		(		_lowerCAmelCase				: Dict    ,				_lowerCAmelCase				: Tuple=0    ,				_lowerCAmelCase				: List[Any]=None						):




 """simple docstring"""







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



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



def     lowerCAmelCase		(		_lowerCAmelCase				: Union[str, Any]    ,				_lowerCAmelCase				: Dict    ,				_lowerCAmelCase				: List[Any]=1    ,				_lowerCAmelCase				: Union[str, Any]=100						):




 """simple docstring"""







 UpperCAmelCase__			,			UpperCAmelCase__							=    jax.lax.top_k(_lowerCAmelCase    ,				_lowerCAmelCase						)
 UpperCAmelCase__			,			UpperCAmelCase__							=    jax.lax.top_k(_lowerCAmelCase    ,				_lowerCAmelCase						)

 UpperCAmelCase__							=    best_end_idx[:, None] - best_start_idx[None, :]
 UpperCAmelCase__							=    jnp.logical_or(widths < 0    ,				widths > max_size						)
 UpperCAmelCase__							=    (best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
 UpperCAmelCase__							=    jnp.argmax(_lowerCAmelCase						).item()

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



def     lowerCAmelCase		(		_lowerCAmelCase				: List[str]						):




 """simple docstring"""







 UpperCAmelCase__							=    sample["question"]["text"]
 UpperCAmelCase__							=    sample["document"]["tokens"]["token"]
 UpperCAmelCase__							=    sample["document"]["tokens"]["is_html"]
 UpperCAmelCase__							=    sample["annotations"]["long_answer"]
 UpperCAmelCase__							=    sample["annotations"]["short_answers"]

 UpperCAmelCase__							=    " ".join([context[i] for i in range(len(_lowerCAmelCase						)						) 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__							=    []
 for s in short_answers:
  short_targets.extend(s["text"]						)
 UpperCAmelCase__							=    list(set(_lowerCAmelCase						)						)

 UpperCAmelCase__							=    []
 for s in long_answers:
  if s["start_token"] == -1:
   continue
  UpperCAmelCase__							=    context[s["start_token"] : s["end_token"]]
  UpperCAmelCase__							=    is_html[s["start_token"] : s["end_token"]]
  UpperCAmelCase__							=    " ".join([answer[i] for i in range(len(_lowerCAmelCase						)						) if not html[i]]						)
  if new_answer not in long_targets:
   long_targets.append(_lowerCAmelCase						)

 UpperCAmelCase__							=    "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     lowerCAmelCase		(		):




 """simple docstring"""







 UpperCAmelCase__							=    load_from_disk("natural-questions-validation"						)
 UpperCAmelCase__							=    dataset.map(_lowerCAmelCase						).remove_columns(["annotations", "document", "id"]						)
 print(_lowerCAmelCase						)

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

 UpperCAmelCase__							=    "vasudevgupta/flax-bigbird-natural-questions"
 UpperCAmelCase__							=    FlaxBigBirdForNaturalQuestions.from_pretrained(_lowerCAmelCase						)
 UpperCAmelCase__							=    BigBirdTokenizerFast.from_pretrained(_lowerCAmelCase						)

 @jax.jit
 def forward(*_lowerCAmelCase				: List[Any]    ,				**_lowerCAmelCase				: Tuple						):
  UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    model(*_lowerCAmelCase    ,				**_lowerCAmelCase						)
  return start_logits, end_logits, jnp.argmax(_lowerCAmelCase    ,				axis=-1						)

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

  UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    forward(**_lowerCAmelCase						)

  UpperCAmelCase__							=    CATEGORY_MAPPING[category.item()]

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

  if predicted_category in ["yes", "no", "null"]:
   UpperCAmelCase__							=    [predicted_category]
   UpperCAmelCase__							=    example["output"] == example["targets"]
   UpperCAmelCase__							=    predicted_category != "null"
   return example

  UpperCAmelCase__							=    38 if predicted_category == "short" else 1024
  UpperCAmelCase__			,			UpperCAmelCase__							=    get_best_valid_start_end_idx(
      start_scores[0]    ,				end_scores[0]    ,				top_k=8    ,				max_size=_lowerCAmelCase						)

  UpperCAmelCase__							=    inputs["input_ids"][0].tolist()
  UpperCAmelCase__							=    [tokenizer.decode(input_ids[start_score : end_score + 1]						)]

  UpperCAmelCase__							=    expand_to_aliases(example["targets"]    ,				make_sub_answers=_lowerCAmelCase						)
  UpperCAmelCase__							=    expand_to_aliases(example["output"]						)

  # some preprocessing to both prediction and answer
  UpperCAmelCase__							=    {"".join(a.split()						) for a in answers}
  UpperCAmelCase__							=    {"".join(p.split()						) for p in predictions}
  UpperCAmelCase__							=    {s for s in predictions if s not in ["``", "''", "`", "'"]}

  # if there is a common element, it's a exact match
  UpperCAmelCase__							=    len(list(answers & predictions						)						) > 0
  UpperCAmelCase__							=    predicted_category != "null" and len(_lowerCAmelCase						) > 0

  return example

 UpperCAmelCase__							=    short_validation_dataset.map(_lowerCAmelCase						)

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


if __name__ == "__main__":
  main()











from math import cos, sin



def     lowerCAmelCase		(		_lowerCAmelCase				: float						):




 """simple docstring"""







 UpperCAmelCase__							=    float(_lowerCAmelCase						)
 return [[scaling_factor * int(x == y						) for x in range(2						)] for y in range(2						)]



def     lowerCAmelCase		(		_lowerCAmelCase				: float						):




 """simple docstring"""







 UpperCAmelCase__			,			UpperCAmelCase__							=    cos(_lowerCAmelCase						), sin(_lowerCAmelCase						)
 return [[c, -s], [s, c]]



def     lowerCAmelCase		(		_lowerCAmelCase				: float						):




 """simple docstring"""







 UpperCAmelCase__			,			UpperCAmelCase__							=    cos(_lowerCAmelCase						), sin(_lowerCAmelCase						)
 UpperCAmelCase__							=    c * s
 return [[c * c, cs], [cs, s * s]]



def     lowerCAmelCase		(		_lowerCAmelCase				: float						):




 """simple docstring"""







 UpperCAmelCase__			,			UpperCAmelCase__							=    cos(_lowerCAmelCase						), sin(_lowerCAmelCase						)
 UpperCAmelCase__							=    c * s
 return [[2 * c - 1, 2 * cs], [2 * cs, 2 * s - 1]]


print(F'''    {scaling(5) = }''')
print(F'''  {rotation(4_5) = }''')
print(F'''{projection(4_5) = }''')


print(F'''{reflection(4_5) = }''')





