






import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				lowercase__											=       42
				lowercase__											=       jnp.floataa
				lowercase__											=       True




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								super().setup()
								lowercase__    =			nn.Dense(5,  dtype=self.dtype  )




				def __call__(     self						:    List[Any],  *lowerCamelCase						:    Dict,  **lowerCamelCase						:    Dict  ):

								'''simple docstring'''
								lowercase__    =			super().__call__(*lowerCamelCase,  **lowerCamelCase  )
								lowercase__    =			self.cls(outputs[2]  )
								return outputs[:2] + (cls_out,)



class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				lowercase__											=       FlaxBigBirdForNaturalQuestionsModule






def   a			(						lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_				):






				'''simple docstring'''
				def cross_entropy(lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_=None				):
								lowercase__    =			logits.shape[-1]
								lowercase__    =			(labels[..., None] == jnp.arange(lowerCamelCase_				)[None]).astype('''f4'''				)
								lowercase__    =			jax.nn.log_softmax(lowerCamelCase_		,   axis=-1				)
								lowercase__    =			-jnp.sum(labels * logits		,   axis=-1				)
								if reduction is not None:
												lowercase__    =			reduction(lowerCamelCase_				)
								return loss

				lowercase__    =			partial(lowerCamelCase_		,   reduction=jnp.mean				)
				lowercase__    =			cross_entropy(lowerCamelCase_		,   lowerCamelCase_				)
				lowercase__    =			cross_entropy(lowerCamelCase_		,   lowerCamelCase_				)
				lowercase__    =			cross_entropy(lowerCamelCase_		,   lowerCamelCase_				)
				return (start_loss + end_loss + pooled_loss) / 3
@dataclass
class       _UpperCAmelCase					:

				"""simple docstring"""




				lowercase__											=       "google/bigbird-roberta-base"
				lowercase__											=       3_000
				lowercase__											=       10_500

				lowercase__											=       128
				lowercase__											=       3

				lowercase__											=       1
				lowercase__											=       5

				# tx_args
				lowercase__											=       3E-5
				lowercase__											=       0.0
				lowercase__											=       20_000
				lowercase__											=       0.0095

				lowercase__											=       "bigbird-roberta-natural-questions"
				lowercase__											=       "training-expt"
				lowercase__											=       "data/nq-training.jsonl"
				lowercase__											=       "data/nq-validation.jsonl"




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								os.makedirs(self.base_dir,  exist_ok=lowerCamelCase  )
								lowercase__    =			os.path.join(self.base_dir,  self.save_dir  )
								lowercase__    =			self.batch_size_per_device * jax.device_count()



@dataclass
class       _UpperCAmelCase					:

				"""simple docstring"""




				lowercase__											=       42
				lowercase__											=       4_096  # no dynamic padding on TPUs




				def __call__(     self						:    List[Any],  lowerCamelCase						:    Optional[Any]  ):

								'''simple docstring'''
								lowercase__    =			self.collate_fn(lowerCamelCase  )
								lowercase__    =			jax.tree_util.tree_map(lowerCamelCase,  lowerCamelCase  )
								return batch




				def 				lowercase__   (     self						:    List[Any],  lowerCamelCase						:    List[str]  ):

								'''simple docstring'''
								lowercase__ ,						lowercase__    =			self.fetch_inputs(features['''input_ids''']  )
								lowercase__    =			{
								    '''input_ids''': jnp.array(lowerCamelCase,  dtype=jnp.intaa  ),
								    '''attention_mask''': jnp.array(lowerCamelCase,  dtype=jnp.intaa  ),
								    '''start_labels''': jnp.array(features['''start_token'''],  dtype=jnp.intaa  ),
								    '''end_labels''': jnp.array(features['''end_token'''],  dtype=jnp.intaa  ),
								    '''pooled_labels''': jnp.array(features['''category'''],  dtype=jnp.intaa  ),
								}
								return batch




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    list  ):

								'''simple docstring'''
								lowercase__    =			[self._fetch_inputs(lowerCamelCase  ) for ids in input_ids]
								return zip(*lowerCamelCase  )




				def 				lowercase__   (     self						:    List[str],  lowerCamelCase						:    list  ):

								'''simple docstring'''
								lowercase__    =			[1 for _ in range(len(lowerCamelCase  )  )]
								while len(lowerCamelCase  ) < self.max_length:
												input_ids.append(self.pad_id  )
												attention_mask.append(0  )
								return input_ids, attention_mask






def   a			(						lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_=None				):






				'''simple docstring'''
				if seed is not None:
								lowercase__    =			dataset.shuffle(seed=lowerCamelCase_				)
				for i in range(len(lowerCamelCase_				) // batch_size				):
								lowercase__    =			dataset[i * batch_size : (i + 1) * batch_size]
								yield dict(lowerCamelCase_				)






@partial(jax.pmap		,   axis_name='''batch'''				)
def   a			(						lowerCamelCase_		,   lowerCamelCase_		,   **lowerCamelCase_				):






				'''simple docstring'''
				def loss_fn(lowerCamelCase_				):
								lowercase__    =			model_inputs.pop('''start_labels'''				)
								lowercase__    =			model_inputs.pop('''end_labels'''				)
								lowercase__    =			model_inputs.pop('''pooled_labels'''				)

								lowercase__    =			state.apply_fn(**lowerCamelCase_		,   params=lowerCamelCase_		,   dropout_rng=lowerCamelCase_		,   train=lowerCamelCase_				)
								lowercase__ ,						lowercase__ ,						lowercase__    =			outputs

								return state.loss_fn(
								    lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   )

				lowercase__ ,						lowercase__    =			jax.random.split(lowerCamelCase_				)
				lowercase__    =			jax.value_and_grad(lowerCamelCase_				)
				lowercase__ ,						lowercase__    =			grad_fn(state.params				)
				lowercase__    =			jax.lax.pmean({'''loss''': loss}		,   axis_name='''batch'''				)
				lowercase__    =			jax.lax.pmean(lowerCamelCase_		,   '''batch'''				)

				lowercase__    =			state.apply_gradients(grads=lowerCamelCase_				)
				return state, metrics, new_drp_rng






@partial(jax.pmap		,   axis_name='''batch'''				)
def   a			(						lowerCamelCase_		,   **lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			model_inputs.pop('''start_labels'''				)
				lowercase__    =			model_inputs.pop('''end_labels'''				)
				lowercase__    =			model_inputs.pop('''pooled_labels'''				)

				lowercase__    =			state.apply_fn(**lowerCamelCase_		,   params=state.params		,   train=lowerCamelCase_				)
				lowercase__ ,						lowercase__ ,						lowercase__    =			outputs

				lowercase__    =			state.loss_fn(lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_				)
				lowercase__    =			jax.lax.pmean({'''loss''': loss}		,   axis_name='''batch'''				)
				return metrics
class       _UpperCAmelCase					(       train_state.TrainState	):

				"""simple docstring"""




				lowercase__											=       struct.field(pytree_node=A__	)



@dataclass
class       _UpperCAmelCase					:

				"""simple docstring"""




				lowercase__											=       42
				lowercase__											=       42
				lowercase__											=       42
				lowercase__											=       42
				lowercase__											=       42
				lowercase__											=       42
				lowercase__											=       None




				def 				lowercase__   (     self						:    List[str],  lowerCamelCase						:    Dict,  lowerCamelCase						:    str,  lowerCamelCase						:    Dict,  lowerCamelCase						:    Any=None  ):

								'''simple docstring'''
								lowercase__    =			model.params
								lowercase__    =			TrainState.create(
								    apply_fn=model.__call__,  params=lowerCamelCase,  tx=lowerCamelCase,  loss_fn=lowerCamelCase,  )
								if ckpt_dir is not None:
												lowercase__ ,						lowercase__ ,						lowercase__ ,						lowercase__ ,						lowercase__    =			restore_checkpoint(lowerCamelCase,  lowerCamelCase  )
												lowercase__    =			{
												    '''lr''': args.lr,
												    '''init_lr''': args.init_lr,
												    '''warmup_steps''': args.warmup_steps,
												    '''num_train_steps''': num_train_steps,
												    '''weight_decay''': args.weight_decay,
												}
												lowercase__ ,						lowercase__    =			build_tx(**lowerCamelCase  )
												lowercase__    =			train_state.TrainState(
												    step=lowerCamelCase,  apply_fn=model.__call__,  params=lowerCamelCase,  tx=lowerCamelCase,  opt_state=lowerCamelCase,  )
												lowercase__    =			args
												lowercase__    =			data_collator
												lowercase__    =			lr
												lowercase__    =			params
								lowercase__    =			jax_utils.replicate(lowerCamelCase  )
								return state




				def 				lowercase__   (     self						:    List[str],  lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Optional[int],  lowerCamelCase						:    Optional[Any]  ):

								'''simple docstring'''
								lowercase__    =			self.args
								lowercase__    =			len(lowerCamelCase  ) // args.batch_size

								lowercase__    =			jax.random.PRNGKey(0  )
								lowercase__    =			jax.random.split(lowerCamelCase,  jax.device_count()  )
								for epoch in range(args.max_epochs  ):
												lowercase__    =			jnp.array(0,  dtype=jnp.floataa  )
												lowercase__    =			get_batched_dataset(lowerCamelCase,  args.batch_size,  seed=lowerCamelCase  )
												lowercase__    =			0
												for batch in tqdm(lowerCamelCase,  total=lowerCamelCase,  desc=F"""Running EPOCH-{epoch}"""  ):
																lowercase__    =			self.data_collator(lowerCamelCase  )
																lowercase__ ,						lowercase__ ,						lowercase__    =			self.train_step_fn(lowerCamelCase,  lowerCamelCase,  **lowerCamelCase  )
																running_loss += jax_utils.unreplicate(metrics['''loss''']  )
																i += 1
																if i % args.logging_steps == 0:
																				lowercase__    =			jax_utils.unreplicate(state.step  )
																				lowercase__    =			running_loss.item() / i
																				lowercase__    =			self.scheduler_fn(state_step - 1  )

																				lowercase__    =			self.evaluate(lowerCamelCase,  lowerCamelCase  )
																				lowercase__    =			{
																				    '''step''': state_step.item(),
																				    '''eval_loss''': eval_loss.item(),
																				    '''tr_loss''': tr_loss,
																				    '''lr''': lr.item(),
																				}
																				tqdm.write(str(lowerCamelCase  )  )
																				self.logger.log(lowerCamelCase,  commit=lowerCamelCase  )

																if i % args.save_steps == 0:
																				self.save_checkpoint(args.save_dir + F"""-e{epoch}-s{i}""",  state=lowerCamelCase  )




				def 				lowercase__   (     self						:    str,  lowerCamelCase						:    Dict,  lowerCamelCase						:    Any  ):

								'''simple docstring'''
								lowercase__    =			get_batched_dataset(lowerCamelCase,  self.args.batch_size  )
								lowercase__    =			len(lowerCamelCase  ) // self.args.batch_size
								lowercase__    =			jnp.array(0,  dtype=jnp.floataa  )
								lowercase__    =			0
								for batch in tqdm(lowerCamelCase,  total=lowerCamelCase,  desc='''Evaluating ... '''  ):
												lowercase__    =			self.data_collator(lowerCamelCase  )
												lowercase__    =			self.val_step_fn(lowerCamelCase,  **lowerCamelCase  )
												running_loss += jax_utils.unreplicate(metrics['''loss''']  )
												i += 1
								return running_loss / i




				def 				lowercase__   (     self						:    Any,  lowerCamelCase						:    Any,  lowerCamelCase						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			jax_utils.unreplicate(lowerCamelCase  )
								print(F"""SAVING CHECKPOINT IN {save_dir}""",  end=''' ... '''  )
								self.model_save_fn(lowerCamelCase,  params=state.params  )
								with open(os.path.join(lowerCamelCase,  '''opt_state.msgpack'''  ),  '''wb'''  ) as f:
												f.write(to_bytes(state.opt_state  )  )
								joblib.dump(self.args,  os.path.join(lowerCamelCase,  '''args.joblib'''  )  )
								joblib.dump(self.data_collator,  os.path.join(lowerCamelCase,  '''data_collator.joblib'''  )  )
								with open(os.path.join(lowerCamelCase,  '''training_state.json'''  ),  '''w'''  ) as f:
												json.dump({'''step''': state.step.item()},  lowerCamelCase  )
								print('''DONE'''  )






def   a			(						lowerCamelCase_		,   lowerCamelCase_				):






				'''simple docstring'''
				print(F"""RESTORING CHECKPOINT FROM {save_dir}"""		,   end=''' ... '''				)
				with open(os.path.join(lowerCamelCase_		,   '''flax_model.msgpack'''				)		,   '''rb'''				) as f:
								lowercase__    =			from_bytes(state.params		,   f.read()				)

				with open(os.path.join(lowerCamelCase_		,   '''opt_state.msgpack'''				)		,   '''rb'''				) as f:
								lowercase__    =			from_bytes(state.opt_state		,   f.read()				)

				lowercase__    =			joblib.load(os.path.join(lowerCamelCase_		,   '''args.joblib'''				)				)
				lowercase__    =			joblib.load(os.path.join(lowerCamelCase_		,   '''data_collator.joblib'''				)				)

				with open(os.path.join(lowerCamelCase_		,   '''training_state.json'''				)		,   '''r'''				) as f:
								lowercase__    =			json.load(lowerCamelCase_				)
				lowercase__    =			training_state['''step''']

				print('''DONE'''				)
				return params, opt_state, step, args, data_collator






def   a			(						lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			num_train_steps - warmup_steps
				lowercase__    =			optax.linear_schedule(init_value=lowerCamelCase_		,   end_value=lowerCamelCase_		,   transition_steps=lowerCamelCase_				)
				lowercase__    =			optax.linear_schedule(init_value=lowerCamelCase_		,   end_value=1e-7		,   transition_steps=lowerCamelCase_				)
				lowercase__    =			optax.join_schedules(schedules=[warmup_fn, decay_fn]		,   boundaries=[warmup_steps]				)
				return lr






def   a			(						lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_				):






				'''simple docstring'''
				def weight_decay_mask(lowerCamelCase_				):
								lowercase__    =			traverse_util.flatten_dict(lowerCamelCase_				)
								lowercase__    =			{k: (v[-1] != '''bias''' and v[-2:] != ('''LayerNorm''', '''scale''')) for k, v in params.items()}
								return traverse_util.unflatten_dict(lowerCamelCase_				)

				lowercase__    =			scheduler_fn(lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_				)

				lowercase__    =			optax.adamw(learning_rate=lowerCamelCase_		,   weight_decay=lowerCamelCase_		,   mask=lowerCamelCase_				)
				return tx, lr












import math
import time

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm
							import torch_xla.debug.metrics as met
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				def __init__(     self						:    Dict,  *lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    Any=None,  **lowerCamelCase						:    str  ):

								'''simple docstring'''
								super().__init__(*lowerCamelCase,  **lowerCamelCase  )
								lowercase__    =			eval_examples
								lowercase__    =			post_process_function




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    str=None,  lowerCamelCase						:    Optional[Any]=None,  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    str = "eval"  ):

								'''simple docstring'''
								lowercase__    =			self.eval_dataset if eval_dataset is None else eval_dataset
								lowercase__    =			self.get_eval_dataloader(lowerCamelCase  )
								lowercase__    =			self.eval_examples if eval_examples is None else eval_examples

								# Temporarily disable metric computation, we will do it in the loop here.
								lowercase__    =			self.compute_metrics
								lowercase__    =			None
								lowercase__    =			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								lowercase__    =			time.time()
								try:
												lowercase__    =			eval_loop(
												    lowerCamelCase,  description='''Evaluation''',  prediction_loss_only=True if compute_metrics is None else None,  ignore_keys=lowerCamelCase,  metric_key_prefix=lowerCamelCase,  )
								finally:
												lowercase__    =			compute_metrics
								lowercase__    =			self.args.eval_batch_size * self.args.world_size
								if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
												start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
								output.metrics.update(
								    speed_metrics(
								        lowerCamelCase,  lowerCamelCase,  num_samples=output.num_samples,  num_steps=math.ceil(output.num_samples / total_batch_size  ),  )  )
								if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
												# Only the main node write the results by default
												lowercase__    =			self.post_process_function(lowerCamelCase,  lowerCamelCase,  output.predictions  )
												lowercase__    =			self.compute_metrics(lowerCamelCase  )

												# Prefix all keys with metric_key_prefix + '_'
												for key in list(metrics.keys()  ):
																if not key.startswith(F"""{metric_key_prefix}_"""  ):
																				lowercase__    =			metrics.pop(lowerCamelCase  )
												metrics.update(output.metrics  )
								else:
												lowercase__    =			output.metrics

								if self.args.should_log:
												# Only the main node log the results by default
												self.log(lowerCamelCase  )

								if self.args.tpu_metrics_debug or self.args.debug:
												# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
												xm.master_print(met.metrics_report()  )

								lowercase__    =			self.callback_handler.on_evaluate(self.args,  self.state,  self.control,  lowerCamelCase  )
								return metrics




				def 				lowercase__   (     self						:    List[Any],  lowerCamelCase						:    Any,  lowerCamelCase						:    Dict,  lowerCamelCase						:    int=None,  lowerCamelCase						:    str = "test"  ):

								'''simple docstring'''
								lowercase__    =			self.get_test_dataloader(lowerCamelCase  )

								# Temporarily disable metric computation, we will do it in the loop here.
								lowercase__    =			self.compute_metrics
								lowercase__    =			None
								lowercase__    =			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								lowercase__    =			time.time()
								try:
												lowercase__    =			eval_loop(
												    lowerCamelCase,  description='''Prediction''',  prediction_loss_only=True if compute_metrics is None else None,  ignore_keys=lowerCamelCase,  metric_key_prefix=lowerCamelCase,  )
								finally:
												lowercase__    =			compute_metrics
								lowercase__    =			self.args.eval_batch_size * self.args.world_size
								if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
												start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
								output.metrics.update(
								    speed_metrics(
								        lowerCamelCase,  lowerCamelCase,  num_samples=output.num_samples,  num_steps=math.ceil(output.num_samples / total_batch_size  ),  )  )

								if self.post_process_function is None or self.compute_metrics is None:
												return output

								lowercase__    =			self.post_process_function(lowerCamelCase,  lowerCamelCase,  output.predictions,  '''predict'''  )
								lowercase__    =			self.compute_metrics(lowerCamelCase  )

								# Prefix all keys with metric_key_prefix + '_'
								for key in list(metrics.keys()  ):
												if not key.startswith(F"""{metric_key_prefix}_"""  ):
																lowercase__    =			metrics.pop(lowerCamelCase  )
								metrics.update(output.metrics  )
								return PredictionOutput(predictions=predictions.predictions,  label_ids=predictions.label_ids,  metrics=lowerCamelCase  )




