



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    (						_UpperCAmelCase     ):
	lowerCAmelCase	:   BigBirdConfig
	lowerCAmelCase	:   jnp.dtype				=	jnp.floataa
	lowerCAmelCase	:   bool				=	True



	def 			__A							(       self		):
							super().setup()
							A__							=       nn.Dense(5			,	dtype=self.dtype		)



	def __call__(       self			,	*UpperCAmelCase__			,	**UpperCAmelCase__		):
							A__							=       super().__call__(*UpperCAmelCase__			,	**UpperCAmelCase__		)
							A__							=       self.cls(outputs[2]		)
							return outputs[:2] + (cls_out,)






class      UpperCamelCase    (						_UpperCAmelCase     ):
	lowerCAmelCase	:   List[Any]				=	FlaxBigBirdForNaturalQuestionsModule





def  UpperCamelCase						(       _A			:				Tuple			,							_A			:				Union[str, Any]			,							_A			:				Optional[Any]			,							_A			:				List[str]			,							_A			:				Optional[int]			,							_A			:				str    )->       Optional[Any]:





						"""simple docstring"""





						def cross_entropy(_A			:				List[Any]			,							_A			:				int			,							_A			:				List[str]=None    ):
												A__							=       logits.shape[-1]
												A__							=       (labels[..., None] == jnp.arange(_A    )[None]).astype("f4"    )
												A__							=       jax.nn.log_softmax(_A			,							axis=-1    )
												A__							=       -jnp.sum(labels * logits			,							axis=-1    )
												if reduction is not None:
																		A__							=       reduction(_A    )
												return loss

						A__							=       partial(_A			,							reduction=jnp.mean    )
						A__							=       cross_entropy(_A			,							_A    )
						A__							=       cross_entropy(_A			,							_A    )
						A__							=       cross_entropy(_A			,							_A    )
						return (start_loss + end_loss + pooled_loss) / 3






@dataclass
class      UpperCamelCase    :
	lowerCAmelCase	:   str				=	"google/bigbird-roberta-base"
	lowerCAmelCase	:   int				=	3_000
	lowerCAmelCase	:   int				=	10_500

	lowerCAmelCase	:   int				=	128
	lowerCAmelCase	:   int				=	3

	lowerCAmelCase	:   int				=	1
	lowerCAmelCase	:   int				=	5

	# tx_args
	lowerCAmelCase	:   float				=	3E-5
	lowerCAmelCase	:   float				=	0.0
	lowerCAmelCase	:   int				=	20_000
	lowerCAmelCase	:   float				=	0.0095

	lowerCAmelCase	:   str				=	"bigbird-roberta-natural-questions"
	lowerCAmelCase	:   str				=	"training-expt"
	lowerCAmelCase	:   str				=	"data/nq-training.jsonl"
	lowerCAmelCase	:   str				=	"data/nq-validation.jsonl"



	def 			__A							(       self		):
							os.makedirs(self.base_dir			,	exist_ok=UpperCAmelCase__		)
							A__							=       os.path.join(self.base_dir			,	self.save_dir		)
							A__							=       self.batch_size_per_device * jax.device_count()






@dataclass
class      UpperCamelCase    :
	lowerCAmelCase	:   int
	lowerCAmelCase	:   int				=	4_096  # no dynamic padding on TPUs



	def __call__(       self			,	UpperCAmelCase__		):
							A__							=       self.collate_fn(UpperCAmelCase__		)
							A__							=       jax.tree_util.tree_map(UpperCAmelCase__			,	UpperCAmelCase__		)
							return batch



	def 			__A							(       self			,	UpperCAmelCase__		):
							A__						,       A__							=       self.fetch_inputs(features["input_ids"]		)
							A__							=       {
							    "input_ids": jnp.array(UpperCAmelCase__			,	dtype=jnp.intaa		),
							    "attention_mask": jnp.array(UpperCAmelCase__			,	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 			__A							(       self			,	UpperCAmelCase__		):
							A__							=       [self._fetch_inputs(UpperCAmelCase__		) for ids in input_ids]
							return zip(*UpperCAmelCase__		)



	def 			__A							(       self			,	UpperCAmelCase__		):
							A__							=       [1 for _ in range(len(UpperCAmelCase__		)		)]
							while len(UpperCAmelCase__		) < self.max_length:
													input_ids.append(self.pad_id		)
													attention_mask.append(0		)
							return input_ids, attention_mask





def  UpperCamelCase						(       _A			:				Optional[Any]			,							_A			:				List[str]			,							_A			:				List[str]=None    )->       Optional[Any]:





						"""simple docstring"""





						if seed is not None:
												A__							=       dataset.shuffle(seed=_A    )
						for i in range(len(_A    ) // batch_size    ):
												A__							=       dataset[i * batch_size : (i + 1) * batch_size]
												yield dict(_A    )
@partial(jax.pmap			,							axis_name="batch"    )
def  UpperCamelCase						(       _A			:				Dict			,							_A			:				Dict			,							**_A			:				str    )->       Dict:





						"""simple docstring"""





						def loss_fn(_A			:				List[Any]    ):
												A__							=       model_inputs.pop("start_labels"    )
												A__							=       model_inputs.pop("end_labels"    )
												A__							=       model_inputs.pop("pooled_labels"    )

												A__							=       state.apply_fn(**_A			,							params=_A			,							dropout_rng=_A			,							train=_A    )
												A__						,       A__						,       A__							=       outputs

												return state.loss_fn(
												    _A			,							_A			,							_A			,							_A			,							_A			,							_A			,							)

						A__						,       A__							=       jax.random.split(_A    )
						A__							=       jax.value_and_grad(_A    )
						A__						,       A__							=       grad_fn(state.params    )
						A__							=       jax.lax.pmean({"loss": loss}			,							axis_name="batch"    )
						A__							=       jax.lax.pmean(_A			,							"batch"    )

						A__							=       state.apply_gradients(grads=_A    )
						return state, metrics, new_drp_rng
@partial(jax.pmap			,							axis_name="batch"    )
def  UpperCamelCase						(       _A			:				List[Any]			,							**_A			:				Dict    )->       str:





						"""simple docstring"""





						A__							=       model_inputs.pop("start_labels"    )
						A__							=       model_inputs.pop("end_labels"    )
						A__							=       model_inputs.pop("pooled_labels"    )

						A__							=       state.apply_fn(**_A			,							params=state.params			,							train=_A    )
						A__						,       A__						,       A__							=       outputs

						A__							=       state.loss_fn(_A			,							_A			,							_A			,							_A			,							_A			,							_A    )
						A__							=       jax.lax.pmean({"loss": loss}			,							axis_name="batch"    )
						return metrics






class      UpperCamelCase    (						train_state.TrainState     ):
	lowerCAmelCase	:   Callable				=	struct.field(pytree_node=_UpperCAmelCase     )






@dataclass
class      UpperCamelCase    :
	lowerCAmelCase	:   Args
	lowerCAmelCase	:   Callable
	lowerCAmelCase	:   Callable
	lowerCAmelCase	:   Callable
	lowerCAmelCase	:   Callable
	lowerCAmelCase	:   wandb
	lowerCAmelCase	:   Callable				=	None



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__=None		):
							A__							=       model.params
							A__							=       TrainState.create(
							    apply_fn=model.__call__			,	params=UpperCAmelCase__			,	tx=UpperCAmelCase__			,	loss_fn=UpperCAmelCase__			,	)
							if ckpt_dir is not None:
													A__						,       A__						,       A__						,       A__						,       A__							=       restore_checkpoint(UpperCAmelCase__			,	UpperCAmelCase__		)
													A__							=       {
													    "lr": args.lr,
													    "init_lr": args.init_lr,
													    "warmup_steps": args.warmup_steps,
													    "num_train_steps": num_train_steps,
													    "weight_decay": args.weight_decay,
													}
													A__						,       A__							=       build_tx(**UpperCAmelCase__		)
													A__							=       train_state.TrainState(
													    step=UpperCAmelCase__			,	apply_fn=model.__call__			,	params=UpperCAmelCase__			,	tx=UpperCAmelCase__			,	opt_state=UpperCAmelCase__			,	)
													A__							=       args
													A__							=       data_collator
													A__							=       lr
													A__							=       params
							A__							=       jax_utils.replicate(UpperCAmelCase__		)
							return state



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       self.args
							A__							=       len(UpperCAmelCase__		) // args.batch_size

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

																									A__							=       self.evaluate(UpperCAmelCase__			,	UpperCAmelCase__		)
																									A__							=       {
																									    "step": state_step.item(),
																									    "eval_loss": eval_loss.item(),
																									    "tr_loss": tr_loss,
																									    "lr": lr.item(),
																									}
																									tqdm.write(str(UpperCAmelCase__		)		)
																									self.logger.log(UpperCAmelCase__			,	commit=UpperCAmelCase__		)

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



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       get_batched_dataset(UpperCAmelCase__			,	self.args.batch_size		)
							A__							=       len(UpperCAmelCase__		) // self.args.batch_size
							A__							=       jnp.array(0			,	dtype=jnp.floataa		)
							A__							=       0
							for batch in tqdm(UpperCAmelCase__			,	total=UpperCAmelCase__			,	desc="Evaluating ... "		):
													A__							=       self.data_collator(UpperCAmelCase__		)
													A__							=       self.val_step_fn(UpperCAmelCase__			,	**UpperCAmelCase__		)
													running_loss += jax_utils.unreplicate(metrics["loss"]		)
													i += 1
							return running_loss / i



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





def  UpperCamelCase						(       _A			:				List[str]			,							_A			:				Union[str, Any]    )->       Any:





						"""simple docstring"""





						print(f"""RESTORING CHECKPOINT FROM {save_dir}"""			,							end=" ... "    )
						with open(os.path.join(_A			,							"flax_model.msgpack"    )			,							"rb"    ) as f:
												A__							=       from_bytes(state.params			,							f.read()    )

						with open(os.path.join(_A			,							"opt_state.msgpack"    )			,							"rb"    ) as f:
												A__							=       from_bytes(state.opt_state			,							f.read()    )

						A__							=       joblib.load(os.path.join(_A			,							"args.joblib"    )    )
						A__							=       joblib.load(os.path.join(_A			,							"data_collator.joblib"    )    )

						with open(os.path.join(_A			,							"training_state.json"    )			,							"r"    ) as f:
												A__							=       json.load(_A    )
						A__							=       training_state["step"]

						print("DONE"    )
						return params, opt_state, step, args, data_collator
def  UpperCamelCase						(       _A			:				List[str]			,							_A			:				Optional[Any]			,							_A			:				Tuple			,							_A			:				List[Any]    )->       Optional[int]:





						"""simple docstring"""





						A__							=       num_train_steps - warmup_steps
						A__							=       optax.linear_schedule(init_value=_A			,							end_value=_A			,							transition_steps=_A    )
						A__							=       optax.linear_schedule(init_value=_A			,							end_value=1E-7			,							transition_steps=_A    )
						A__							=       optax.join_schedules(schedules=[warmup_fn, decay_fn]			,							boundaries=[warmup_steps]    )
						return lr
def  UpperCamelCase						(       _A			:				Optional[Any]			,							_A			:				str			,							_A			:				Any			,							_A			:				Any			,							_A			:				Optional[Any]    )->       List[Any]:





						"""simple docstring"""





						def weight_decay_mask(_A			:				int    ):
												A__							=       traverse_util.flatten_dict(_A    )
												A__							=       {k: (v[-1] != "bias" and v[-2:] != ("LayerNorm", "scale")) for k, v in params.items()}
												return traverse_util.unflatten_dict(_A    )

						A__							=       scheduler_fn(_A			,							_A			,							_A			,							_A    )

						A__							=       optax.adamw(learning_rate=_A			,							weight_decay=_A			,							mask=_A    )
						return tx, lr










import math
from enum import Enum
from typing import Optional, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from .utils import logging


UpperCAmelCase_							:      Optional[Any]      =  logging.get_logger(__name__)






class      UpperCamelCase    (						_UpperCAmelCase     ):
	lowerCAmelCase	:   List[str]				=	"""linear"""
	lowerCAmelCase	:   int				=	"""cosine"""
	lowerCAmelCase	:   Dict				=	"""cosine_with_restarts"""
	lowerCAmelCase	:   Optional[Any]				=	"""polynomial"""
	lowerCAmelCase	:   Dict				=	"""constant"""
	lowerCAmelCase	:   Any				=	"""constant_with_warmup"""
	lowerCAmelCase	:   Union[str, Any]				=	"""piecewise_constant"""





def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int = -1    )->       Dict:





						"""simple docstring"""





						return LambdaLR(_A			,							lambda _A    : 1			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int = -1    )->       Optional[Any]:





						"""simple docstring"""





						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1.0			,							_A    )    )
												return 1.0

						return LambdaLR(_A			,							_A			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				str			,							_A			:				int = -1    )->       Dict:





						"""simple docstring"""





						A__							=       {}
						A__							=       step_rules.split(","    )
						for rule_str in rule_list[:-1]:
												A__						,       A__							=       rule_str.split(":"    )
												A__							=       int(_A    )
												A__							=       float(_A    )
												A__							=       value
						A__							=       float(rule_list[-1]    )

						def create_rules_function(_A			:				Any			,							_A			:				Optional[int]    ):
												def rule_func(_A			:				int    ) -> float:
																		A__							=       sorted(rules_dict.keys()    )
																		for i, sorted_step in enumerate(_A    ):
																								if steps < sorted_step:
																														return rules_dict[sorted_steps[i]]
																		return last_lr_multiple

												return rule_func

						A__							=       create_rules_function(_A			,							_A    )

						return LambdaLR(_A			,							_A			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Any			,							_A			:				Union[str, Any]			,							_A			:				str			,							_A			:				str=-1    )->       Tuple:





						"""simple docstring"""





						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												return max(
												    0.0			,							float(num_training_steps - current_step    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int			,							_A			:				float = 0.5			,							_A			:				int = -1    )->       Any:





						"""simple docstring"""





						def lr_lambda(_A			:				Tuple    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												A__							=       float(current_step - num_warmup_steps    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )
												return max(0.0			,							0.5 * (1.0 + math.cos(math.pi * float(_A    ) * 2.0 * progress    ))    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int			,							_A			:				int = 1			,							_A			:				int = -1    )->       Any:





						"""simple docstring"""





						def lr_lambda(_A			:				Tuple    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												A__							=       float(current_step - num_warmup_steps    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )
												if progress >= 1.0:
																		return 0.0
												return max(0.0			,							0.5 * (1.0 + math.cos(math.pi * ((float(_A    ) * progress) % 1.0)    ))    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Union[str, Any]			,							_A			:				Union[str, Any]			,							_A			:				List[str]			,							_A			:				Tuple=1E-7			,							_A			:				Dict=1.0			,							_A			:				Union[str, Any]=-1    )->       Any:





						"""simple docstring"""





						A__							=       optimizer.defaults["lr"]
						if not (lr_init > lr_end):
												raise ValueError(f"""lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"""    )

						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												elif current_step > num_training_steps:
																		return lr_end / lr_init  # as LambdaLR multiplies by lr_init
												else:
																		A__							=       lr_init - lr_end
																		A__							=       num_training_steps - num_warmup_steps
																		A__							=       1 - (current_step - num_warmup_steps) / decay_steps
																		A__							=       lr_range * pct_remaining**power + lr_end
																		return decay / lr_init  # as LambdaLR multiplies by lr_init

						return LambdaLR(_A			,							_A			,							_A    )


UpperCAmelCase_							:      Any      =  {
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
}
def  UpperCamelCase						(       _A			:				Union[str, SchedulerType]			,							_A			:				Optimizer			,							_A			:				Optional[str] = None			,							_A			:				Optional[int] = None			,							_A			:				Optional[int] = None			,							_A			:				int = 1			,							_A			:				float = 1.0			,							_A			:				int = -1			,							)->       Union[str, Any]:





						"""simple docstring"""





						A__							=       SchedulerType(_A    )
						A__							=       TYPE_TO_SCHEDULER_FUNCTION[name]
						if name == SchedulerType.CONSTANT:
												return schedule_func(_A			,							last_epoch=_A    )

						if name == SchedulerType.PIECEWISE_CONSTANT:
												return schedule_func(_A			,							step_rules=_A			,							last_epoch=_A    )

						# All other schedulers require `num_warmup_steps`
						if num_warmup_steps is None:
												raise ValueError(f"""{name} requires `num_warmup_steps`, please provide that argument."""    )

						if name == SchedulerType.CONSTANT_WITH_WARMUP:
												return schedule_func(_A			,							num_warmup_steps=_A			,							last_epoch=_A    )

						# All other schedulers require `num_training_steps`
						if num_training_steps is None:
												raise ValueError(f"""{name} requires `num_training_steps`, please provide that argument."""    )

						if name == SchedulerType.COSINE_WITH_RESTARTS:
												return schedule_func(
												    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							num_cycles=_A			,							last_epoch=_A			,							)

						if name == SchedulerType.POLYNOMIAL:
												return schedule_func(
												    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							power=_A			,							last_epoch=_A			,							)

						return schedule_func(
						    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							last_epoch=_A    )





