






'''simple docstring'''


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

from transformers import BigBirdTokenizerFast


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






def   lowerCAmelCase__			(					SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__=0       ,		SCREAMING_SNAKE_CASE__=None		):
 return [" ".join(x.split(' '		)[begin:end]		) for x in answers if len(x.split(' '		)		) > 1]






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

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

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






def   lowerCAmelCase__			(					SCREAMING_SNAKE_CASE__		):
 __a    :	Optional[Any]									=						sample['question']['text']
 __a    :	Optional[int]									=						sample['document']['tokens']['token']
 __a    :	Any									=						sample['document']['tokens']['is_html']
 __a    :	str									=						sample['annotations']['long_answer']
 __a    :	int									=						sample['annotations']['short_answers']

 __a    :	Optional[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",
   }

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

 __a    :	List[Any]									=						[]
 for s in long_answers:
  if s["start_token"] == -1:
   continue
  __a    :	Optional[int]									=						context[s['start_token'] : s['end_token']]
  __a    :	List[Any]									=						is_html[s['start_token'] : s['end_token']]
  __a    :	int									=						' '.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__		)

 __a    :	List[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   lowerCAmelCase__			(					):
 __a    :	List[Any]									=						load_from_disk('natural-questions-validation'		)
 __a    :	Union[str, Any]									=						dataset.map(SCREAMING_SNAKE_CASE__		).remove_columns(['annotations', 'document', 'id']		)
 print(SCREAMING_SNAKE_CASE__		)

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

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

 @jax.jit
 def forward(*SCREAMING_SNAKE_CASE__       ,		**SCREAMING_SNAKE_CASE__		):
  __a    :	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
  __a    :	List[str]									=						tokenizer(
      example['question']       ,		example['context']       ,		return_tensors='np'       ,		max_length=4096       ,		padding='max_length'       ,		truncation=SCREAMING_SNAKE_CASE__       ,		)

  __a    :	Optional[Any]									=						forward(**SCREAMING_SNAKE_CASE__		)

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

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

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

  __a    :	List[str]									=						38 if predicted_category == 'short' else 1024
  __a    :	Optional[int]									=						get_best_valid_start_end_idx(
      start_scores[0]       ,		end_scores[0]       ,		top_k=8       ,		max_size=SCREAMING_SNAKE_CASE__		)

  __a    :	Optional[int]									=						inputs['input_ids'][0].tolist()
  __a    :	Optional[Any]									=						[tokenizer.decode(input_ids[start_score : end_score + 1]		)]

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

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

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

  return example

 __a    :	str									=						short_validation_dataset.map(SCREAMING_SNAKE_CASE__		)

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


if __name__ == "__main__":
  main()













'''simple docstring'''


from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
  import torch

SCREAMING_SNAKE_CASE_				      =							logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
  from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):






  def   lowerCAmelCase__			(					cls		):
   return cls






def   lowerCAmelCase__			(					SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__ = None       ,		SCREAMING_SNAKE_CASE__ = None       ,		SCREAMING_SNAKE_CASE__ = None       ,		SCREAMING_SNAKE_CASE__ = None       ,		):
 __a    :	int									=						device
 __a    :	Union[str, Any]									=						shape[0]

 __a    :	Union[str, Any]									=						layout or torch.strided
 __a    :	Any									=						device or torch.device('cpu'		)

 if generator is not None:
  __a    :	Optional[int]									=						generator.device.type if not isinstance(SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__		) else generator[0].device.type
  if gen_device_type != device.type and gen_device_type == "cpu":
   __a    :	int									=						'cpu'
   if device != "mps":
    logger.info(
        f'''The passed generator was created on \'cpu\' even though a tensor on {device} was expected.'''
        f''' Tensors will be created on \'cpu\' and then moved to {device}. Note that one can probably'''
        f''' slighly speed up this function by passing a generator that was created on the {device} device.'''		)
  elif gen_device_type != device.type and gen_device_type == "cuda":
   raise ValueError(f'''Cannot generate a {device} tensor from a generator of type {gen_device_type}.'''		)

 if isinstance(SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__		):
  __a    :	Optional[int]									=						(1,) + shape[1:]
  __a    :	Dict									=						[
      torch.randn(SCREAMING_SNAKE_CASE__       ,		generator=generator[i]       ,		device=SCREAMING_SNAKE_CASE__       ,		dtype=SCREAMING_SNAKE_CASE__       ,		layout=SCREAMING_SNAKE_CASE__		)
      for i in range(SCREAMING_SNAKE_CASE__		)
  ]
  __a    :	List[str]									=						torch.cat(SCREAMING_SNAKE_CASE__       ,		dim=0		).to(SCREAMING_SNAKE_CASE__		)
 else:
  __a    :	Tuple									=						torch.randn(SCREAMING_SNAKE_CASE__       ,		generator=SCREAMING_SNAKE_CASE__       ,		device=SCREAMING_SNAKE_CASE__       ,		dtype=SCREAMING_SNAKE_CASE__       ,		layout=SCREAMING_SNAKE_CASE__		).to(SCREAMING_SNAKE_CASE__		)

 return latents






def   lowerCAmelCase__			(					SCREAMING_SNAKE_CASE__		):
 if is_torch_version('<'       ,		'2.0.0'		) or not hasattr(SCREAMING_SNAKE_CASE__       ,		'_dynamo'		):
  return False
 return isinstance(SCREAMING_SNAKE_CASE__       ,		torch._dynamo.eval_frame.OptimizedModule		)




