






"""simple docstring"""



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

from transformers import BigBirdTokenizerFast


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




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




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




def     lowerCamelCase_					(	__A	:		str      ,						__A	:		int      ,						__A	:		Dict=1      ,						__A	:		str=100      ):
					lowerCamelCase__					,  lowerCamelCase__					=jax.lax.top_k(__A      ,						__A      )
					lowerCamelCase__					,  lowerCamelCase__					=jax.lax.top_k(__A      ,						__A      )

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

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




def     lowerCamelCase_					(	__A	:		int      ):
					lowerCamelCase__					=sample['question']['text']
					lowerCamelCase__					=sample['document']['tokens']['token']
					lowerCamelCase__					=sample['document']['tokens']['is_html']
					lowerCamelCase__					=sample['annotations']['long_answer']
					lowerCamelCase__					=sample['annotations']['short_answers']

					lowerCamelCase__					=' '.join([context[i] for i in range(len(__A      )      ) 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",
															}

					lowerCamelCase__					=[]
					for s in short_answers:
										short_targets.extend(s['text']      )
					lowerCamelCase__					=list(set(__A      )      )

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

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

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

					lowerCamelCase__					='vasudevgupta/flax-bigbird-natural-questions'
					lowerCamelCase__					=FlaxBigBirdForNaturalQuestions.from_pretrained(__A      )
					lowerCamelCase__					=BigBirdTokenizerFast.from_pretrained(__A      )

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

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

										lowerCamelCase__					,  lowerCamelCase__					,  lowerCamelCase__					=forward(**__A      )

										lowerCamelCase__					=CATEGORY_MAPPING[category.item()]

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

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

										lowerCamelCase__					=38 if predicted_category == 'short' else 1024
										lowerCamelCase__					,  lowerCamelCase__					=get_best_valid_start_end_idx(
										    start_scores[0]      ,						end_scores[0]      ,						top_k=8      ,						max_size=__A      )

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

										lowerCamelCase__					=expand_to_aliases(example['targets']      ,						make_sub_answers=__A      )
										lowerCamelCase__					=expand_to_aliases(example['output']      )

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

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

										return example

					lowerCamelCase__					=short_validation_dataset.map(__A      )

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


if __name__ == "__main__":
						main()







"""simple docstring"""



import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
						import torch

						from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
						import pyanvml.pyanvml as nvml


_snake_case							:   str	  =       logging.get_logger(__name__)




class 							__lowercase(							_UpperCamelCase	):
						lowerCAmelCase_		:				PyTorchBenchmarkArguments
						lowerCAmelCase_		:				PretrainedConfig
						lowerCAmelCase_		:				str=							"PyTorch"






						@property
						def 		lowerCAmelCase (       self)      ->	List[Any]:




											"""simple docstring"""





											return torch.__version__






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_inference_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_speed(_inference)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, Optional[MemorySummary]]:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_inference_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_memory(_inference)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_train_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_speed(_train)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, Optional[MemorySummary]]:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_train_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_memory(_train)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Callable[[], None]:




											"""simple docstring"""





											lowerCamelCase__					=self.config_dict[model_name]

											if self.args.torchscript:
																lowerCamelCase__					=True

											lowerCamelCase__					=(
											    hasattr(_SCREAMING_SNAKE_CASE     ,    'architectures')
											    and isinstance(config.architectures     ,    _SCREAMING_SNAKE_CASE)
											    and len(config.architectures) > 0
											)
											if not self.args.only_pretrain_model and has_model_class_in_config:
																try:
																					lowerCamelCase__					=config.architectures[0]
																					lowerCamelCase__					=__import__('transformers'     ,    fromlist=[model_class])
																					lowerCamelCase__					=getattr(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=model_cls(_SCREAMING_SNAKE_CASE)
																except ImportError:
																					raise ImportError(
																					    F"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
																					    ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.')
											else:
																lowerCamelCase__					=MODEL_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE)

											model.eval()
											model.to(self.args.device)

											# encoder-decoder has vocab size saved differently
											lowerCamelCase__					=config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE     ,    'vocab_size') else config.encoder.vocab_size
											lowerCamelCase__					=torch.randint(_SCREAMING_SNAKE_CASE     ,    (batch_size, sequence_length)     ,    dtype=torch.long     ,    device=self.args.device)

											if self.args.fpaa:
																logger.info('Running training in Mixed Precision...')
																if not self.args.is_gpu:
																					raise ValueError('Mixed precision is possible only for GPU.')
																# amp seems to have memory leaks so that memory usage
																# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
																model.half()

											if self.args.torchscript:
																with torch.no_grad():
																					lowerCamelCase__					=torch.jit.trace(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											else:
																lowerCamelCase__					=model

											def encoder_decoder_forward():
																with torch.no_grad():
																					lowerCamelCase__					=inference_model(_SCREAMING_SNAKE_CASE     ,    decoder_input_ids=_SCREAMING_SNAKE_CASE)
																return outputs

											def encoder_forward():
																with torch.no_grad():
																					lowerCamelCase__					=inference_model(_SCREAMING_SNAKE_CASE)
																return outputs

											lowerCamelCase__					=encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
											return _forward






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Callable[[], None]:




											"""simple docstring"""





											lowerCamelCase__					=self.config_dict[model_name]

											lowerCamelCase__					=(
											    hasattr(_SCREAMING_SNAKE_CASE     ,    'architectures')
											    and isinstance(config.architectures     ,    _SCREAMING_SNAKE_CASE)
											    and len(config.architectures) > 0
											)
											if not self.args.only_pretrain_model and has_model_class_in_config:
																try:
																					lowerCamelCase__					=config.architectures[0]
																					lowerCamelCase__					=__import__('transformers'     ,    fromlist=[model_class])
																					lowerCamelCase__					=getattr(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=model_cls(_SCREAMING_SNAKE_CASE)
																except ImportError:
																					raise ImportError(
																					    F"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
																					    ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.')
											else:
																lowerCamelCase__					=MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE)

											if self.args.torchscript:
																raise NotImplementedError('Training for torchscript is currently not implemented')
											else:
																lowerCamelCase__					=model

											model.train()
											model.to(self.args.device)

											# encoder-decoder has vocab size saved differently
											lowerCamelCase__					=config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE     ,    'vocab_size') else config.encoder.vocab_size
											lowerCamelCase__					=torch.randint(_SCREAMING_SNAKE_CASE     ,    (batch_size, sequence_length)     ,    dtype=torch.long     ,    device=self.args.device)

											if self.args.fpaa:
																logger.info('Running training in Mixed Precision...')
																if not self.args.is_gpu:
																					raise ValueError('Mixed precision is possible only for GPU.')

																# amp seems to have memory leaks so that memory usage
																# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
																model.half()

											def compute_loss_and_backprob_encoder():
																lowerCamelCase__					=train_model(_SCREAMING_SNAKE_CASE     ,    labels=_SCREAMING_SNAKE_CASE)[0]
																loss.backward()
																return loss

											def compute_loss_and_backprob_encoder_decoder():
																lowerCamelCase__					=train_model(_SCREAMING_SNAKE_CASE     ,    decoder_input_ids=_SCREAMING_SNAKE_CASE     ,    labels=_SCREAMING_SNAKE_CASE)[0]
																loss.backward()
																return loss

											lowerCamelCase__					=(
											    compute_loss_and_backprob_encoder_decoder
											    if config.is_encoder_decoder
											    else compute_loss_and_backprob_encoder
											)
											return _train






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""simple docstring"""





											try:
																if self.args.is_tpu or self.args.torchscript:
																					# run additional 10 times to stabilize compilation for tpu and torchscript
																					logger.info('Do inference on TPU or torchscript. Running model 5 times to stabilize compilation')
																					timeit.repeat(
																					    _SCREAMING_SNAKE_CASE     ,    repeat=1     ,    number=5     ,    )

																# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
																lowerCamelCase__					=timeit.repeat(
																    _SCREAMING_SNAKE_CASE     ,    repeat=self.args.repeat     ,    number=10     ,    )

																if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
																					import torch_xla.debug.metrics as met

																					self.print_fn(met.metrics_report())

																return min(_SCREAMING_SNAKE_CASE) / 10.0
											except RuntimeError as e:
																self.print_fn(F"Doesn't fit on GPU. {e}")
																return "N/A"






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, MemorySummary]:




											"""simple docstring"""





											try:
																if self.args.trace_memory_line_by_line:
																					lowerCamelCase__					=start_memory_tracing('transformers')

																if self.args.is_tpu:
																					# tpu
																					raise NotImplementedError(
																					    'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'
																					    ' `--no-memory` or `args.memory=False`')
																elif self.args.is_gpu:
																					if not is_pyanvml_available():
																										logger.warning(
																										    'py3nvml not installed, we won\'t log GPU memory usage. '
																										    'Install py3nvml (pip install py3nvml) to log information about GPU.')
																										lowerCamelCase__					='N/A'
																					else:
																										logger.info(
																										    'Measuring total GPU usage on GPU device. Make sure to not have additional processes running'
																										    ' on the same GPU.')
																										# init nvml
																										nvml.nvmlInit()
																										func()
																										lowerCamelCase__					=nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
																										lowerCamelCase__					=nvml.nvmlDeviceGetMemoryInfo(_SCREAMING_SNAKE_CASE)
																										lowerCamelCase__					=meminfo.used
																										lowerCamelCase__					=Memory(_SCREAMING_SNAKE_CASE)
																										# shutdown nvml
																										nvml.nvmlShutdown()
																else:
																					# cpu
																					lowerCamelCase__					=measure_peak_memory_cpu(_SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=Memory(_SCREAMING_SNAKE_CASE) if isinstance(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE) else memory_bytes

																if self.args.trace_memory_line_by_line:
																					lowerCamelCase__					=stop_memory_tracing(_SCREAMING_SNAKE_CASE)
																else:
																					lowerCamelCase__					=None

																return memory, summary
											except RuntimeError as e:
																self.print_fn(F"Doesn't fit on GPU. {e}")
																return "N/A", None