





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       lowercase					(				lowerCamelCase__  ):



						"""simple docstring"""

						_UpperCamelCase	:						BigBirdConfig
						_UpperCamelCase	:						jnp.dtype    		=						jnp.floataa
						_UpperCamelCase	:						bool    		=						True







						def 		__UpperCAmelCase							(     self      :		List[str]       ):





										'''simple docstring'''







										super().setup()
										_snake_case			:      Dict	    = nn.Dense(5			,					dtype=self.dtype       )







						def __call__(     self      :		Any			,					*lowerCamelCase_      :		Any			,					**lowerCamelCase_      :		Optional[Any]       ):





										'''simple docstring'''







										_snake_case			:      List[Any]	    = super().__call__(*snake_case__			,					**snake_case__       )
										_snake_case			:      int	    = self.cls(outputs[2]       )
										return outputs[:2] + (cls_out,)



class       lowercase					(				lowerCamelCase__  ):



						"""simple docstring"""

						_UpperCamelCase	:						Tuple    		=						FlaxBigBirdForNaturalQuestionsModule



def    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase			):



				def cross_entropy(__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase=None			):
								_snake_case			:      Optional[int]	    = logits.shape[-1]
								_snake_case			:      int	    = (labels[..., None] == jnp.arange(__lowerCAmelCase			)[None]).astype('f4'			)
								_snake_case			:      Tuple	    = jax.nn.log_softmax(__lowerCAmelCase						,      axis=-1			)
								_snake_case			:      List[str]	    = -jnp.sum(labels * logits						,      axis=-1			)
								if reduction is not None:
												_snake_case			:      Dict	    = reduction(__lowerCAmelCase			)
								return loss

				_snake_case			:      Dict	    = partial(__lowerCAmelCase						,      reduction=jnp.mean			)
				_snake_case			:      Optional[int]	    = cross_entropy(__lowerCAmelCase						,      __lowerCAmelCase			)
				_snake_case			:      List[Any]	    = cross_entropy(__lowerCAmelCase						,      __lowerCAmelCase			)
				_snake_case			:      Optional[int]	    = cross_entropy(__lowerCAmelCase						,      __lowerCAmelCase			)
				return (start_loss + end_loss + pooled_loss) / 3


@dataclass
class       lowercase					:



						"""simple docstring"""

						_UpperCamelCase	:						str    		=						"google/bigbird-roberta-base"
						_UpperCamelCase	:						int    		=						3000
						_UpperCamelCase	:						int    		=						1_0500

						_UpperCamelCase	:						int    		=						128
						_UpperCamelCase	:						int    		=						3

						_UpperCamelCase	:						int    		=						1
						_UpperCamelCase	:						int    		=						5

						# tx_args
						_UpperCamelCase	:						float    		=						3e-5
						_UpperCamelCase	:						float    		=						0.0
						_UpperCamelCase	:						int    		=						2_0000
						_UpperCamelCase	:						float    		=						0.00_95

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







						def 		__UpperCAmelCase							(     self      :		List[Any]       ):





										'''simple docstring'''







										os.makedirs(self.base_dir			,					exist_ok=snake_case__       )
										_snake_case			:      str	    = os.path.join(self.base_dir			,					self.save_dir       )
										_snake_case			:      Optional[Any]	    = self.batch_size_per_device * jax.device_count()



@dataclass
class       lowercase					:



						"""simple docstring"""

						_UpperCamelCase	:						int
						_UpperCamelCase	:						int    		=						4096  # no dynamic padding on TPUs







						def __call__(     self      :		Any			,					lowerCamelCase_      :		List[Any]       ):





										'''simple docstring'''







										_snake_case			:      Tuple	    = self.collate_fn(snake_case__       )
										_snake_case			:      Union[str, Any]	    = jax.tree_util.tree_map(snake_case__			,					snake_case__       )
										return batch







						def 		__UpperCAmelCase							(     self      :		Any			,					lowerCamelCase_      :		Optional[Any]       ):





										'''simple docstring'''







										_snake_case			:      Tuple	    = self.fetch_inputs(features['input_ids']       )
										_snake_case			:      List[Any]	    = {
										    """input_ids""": jnp.array(snake_case__			,					dtype=jnp.intaa       ),
										    """attention_mask""": jnp.array(snake_case__			,					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 		__UpperCAmelCase							(     self      :		Optional[int]			,					lowerCamelCase_      :		List[Any]       ):





										'''simple docstring'''







										_snake_case			:      Optional[Any]	    = [self._fetch_inputs(snake_case__       ) for ids in input_ids]
										return zip(*snake_case__       )







						def 		__UpperCAmelCase							(     self      :		List[Any]			,					lowerCamelCase_      :		int       ):





										'''simple docstring'''







										_snake_case			:      List[str]	    = [1 for _ in range(len(snake_case__       )       )]
										while len(snake_case__       ) < self.max_length:
														input_ids.append(self.pad_id       )
														attention_mask.append(0       )
										return input_ids, attention_mask



def    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase=None			):



				if seed is not None:
								_snake_case			:      Optional[Any]	    = dataset.shuffle(seed=__lowerCAmelCase			)
				for i in range(len(__lowerCAmelCase			) // batch_size			):
								_snake_case			:      Any	    = dataset[i * batch_size : (i + 1) * batch_size]
								yield dict(__lowerCAmelCase			)

@partial(jax.pmap						,      axis_name='batch'			)
def    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      **__lowerCAmelCase			):



				def loss_fn(__lowerCAmelCase			):
								_snake_case			:      Union[str, Any]	    = model_inputs.pop('start_labels'			)
								_snake_case			:      int	    = model_inputs.pop('end_labels'			)
								_snake_case			:      str	    = model_inputs.pop('pooled_labels'			)

								_snake_case			:      List[str]	    = state.apply_fn(**__lowerCAmelCase						,      params=__lowerCAmelCase						,      dropout_rng=__lowerCAmelCase						,      train=__lowerCAmelCase			)
								_snake_case			:      Union[str, Any]	    = outputs

								return state.loss_fn(
								    __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      )

				_snake_case			:      Optional[int]	    = jax.random.split(__lowerCAmelCase			)
				_snake_case			:      List[str]	    = jax.value_and_grad(__lowerCAmelCase			)
				_snake_case			:      Optional[Any]	    = grad_fn(state.params			)
				_snake_case			:      List[Any]	    = jax.lax.pmean({'loss': loss}						,      axis_name='batch'			)
				_snake_case			:      Tuple	    = jax.lax.pmean(__lowerCAmelCase						,      'batch'			)

				_snake_case			:      Union[str, Any]	    = state.apply_gradients(grads=__lowerCAmelCase			)
				return state, metrics, new_drp_rng

@partial(jax.pmap						,      axis_name='batch'			)
def    A__(		__lowerCAmelCase						,      **__lowerCAmelCase			):
				_snake_case			:      Union[str, Any]	    = model_inputs.pop('start_labels'			)
				_snake_case			:      int	    = model_inputs.pop('end_labels'			)
				_snake_case			:      List[str]	    = model_inputs.pop('pooled_labels'			)

				_snake_case			:      Optional[int]	    = state.apply_fn(**__lowerCAmelCase						,      params=state.params						,      train=__lowerCAmelCase			)
				_snake_case			:      Optional[int]	    = outputs

				_snake_case			:      List[Any]	    = state.loss_fn(__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase			)
				_snake_case			:      List[Any]	    = jax.lax.pmean({'loss': loss}						,      axis_name='batch'			)
				return metrics


class       lowercase					(				train_state.TrainState  ):



						"""simple docstring"""

						_UpperCamelCase	:						Callable    		=						struct.field(pytree_node=lowerCamelCase__  )



@dataclass
class       lowercase					:



						"""simple docstring"""

						_UpperCamelCase	:						Args
						_UpperCamelCase	:						Callable
						_UpperCamelCase	:						Callable
						_UpperCamelCase	:						Callable
						_UpperCamelCase	:						Callable
						_UpperCamelCase	:						wandb
						_UpperCamelCase	:						Callable    		=						None







						def 		__UpperCAmelCase							(     self      :		Optional[int]			,					lowerCamelCase_      :		Tuple			,					lowerCamelCase_      :		int			,					lowerCamelCase_      :		List[Any]			,					lowerCamelCase_      :		List[str]=None       ):





										'''simple docstring'''







										_snake_case			:      List[Any]	    = model.params
										_snake_case			:      Tuple	    = TrainState.create(
										    apply_fn=model.__call__			,					params=snake_case__			,					tx=snake_case__			,					loss_fn=snake_case__			,					)
										if ckpt_dir is not None:
														_snake_case			:      Union[str, Any]	    = restore_checkpoint(snake_case__			,					snake_case__       )
														_snake_case			:      Optional[Any]	    = {
														    """lr""": args.lr,
														    """init_lr""": args.init_lr,
														    """warmup_steps""": args.warmup_steps,
														    """num_train_steps""": num_train_steps,
														    """weight_decay""": args.weight_decay,
														}
														_snake_case			:      Union[str, Any]	    = build_tx(**snake_case__       )
														_snake_case			:      Optional[Any]	    = train_state.TrainState(
														    step=snake_case__			,					apply_fn=model.__call__			,					params=snake_case__			,					tx=snake_case__			,					opt_state=snake_case__			,					)
														_snake_case			:      str	    = args
														_snake_case			:      Dict	    = data_collator
														_snake_case			:      Tuple	    = lr
														_snake_case			:      Any	    = params
										_snake_case			:      Any	    = jax_utils.replicate(snake_case__       )
										return state







						def 		__UpperCAmelCase							(     self      :		Any			,					lowerCamelCase_      :		int			,					lowerCamelCase_      :		List[str]			,					lowerCamelCase_      :		Tuple       ):





										'''simple docstring'''







										_snake_case			:      Optional[Any]	    = self.args
										_snake_case			:      str	    = len(snake_case__       ) // args.batch_size

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

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

																		if i % args.save_steps == 0:
																						self.save_checkpoint(args.save_dir + f'''-e{epoch}-s{i}'''			,					state=snake_case__       )







						def 		__UpperCAmelCase							(     self      :		List[str]			,					lowerCamelCase_      :		List[str]			,					lowerCamelCase_      :		Any       ):





										'''simple docstring'''







										_snake_case			:      Union[str, Any]	    = get_batched_dataset(snake_case__			,					self.args.batch_size       )
										_snake_case			:      List[str]	    = len(snake_case__       ) // self.args.batch_size
										_snake_case			:      Optional[Any]	    = jnp.array(0			,					dtype=jnp.floataa       )
										_snake_case			:      Optional[int]	    = 0
										for batch in tqdm(snake_case__			,					total=snake_case__			,					desc='Evaluating ... '       ):
														_snake_case			:      Any	    = self.data_collator(snake_case__       )
														_snake_case			:      int	    = self.val_step_fn(snake_case__			,					**snake_case__       )
														running_loss += jax_utils.unreplicate(metrics['loss']       )
														i += 1
										return running_loss / i







						def 		__UpperCAmelCase							(     self      :		List[str]			,					lowerCamelCase_      :		Optional[Any]			,					lowerCamelCase_      :		Tuple       ):





										'''simple docstring'''







										_snake_case			:      Optional[int]	    = jax_utils.unreplicate(snake_case__       )
										print(f'''SAVING CHECKPOINT IN {save_dir}'''			,					end=' ... '       )
										self.model_save_fn(snake_case__			,					params=state.params       )
										with open(os.path.join(snake_case__			,					'opt_state.msgpack'       )			,					'wb'       ) as f:
														f.write(to_bytes(state.opt_state       )       )
										joblib.dump(self.args			,					os.path.join(snake_case__			,					'args.joblib'       )       )
										joblib.dump(self.data_collator			,					os.path.join(snake_case__			,					'data_collator.joblib'       )       )
										with open(os.path.join(snake_case__			,					'training_state.json'       )			,					'w'       ) as f:
														json.dump({'step': state.step.item()}			,					snake_case__       )
										print('DONE'       )



def    A__(		__lowerCAmelCase						,      __lowerCAmelCase			):
				print(F'''RESTORING CHECKPOINT FROM {save_dir}'''						,      end=' ... '			)
				with open(os.path.join(__lowerCAmelCase						,      'flax_model.msgpack'			)						,      'rb'			) as f:
								_snake_case			:      Optional[int]	    = from_bytes(state.params						,      f.read()			)

				with open(os.path.join(__lowerCAmelCase						,      'opt_state.msgpack'			)						,      'rb'			) as f:
								_snake_case			:      Optional[int]	    = from_bytes(state.opt_state						,      f.read()			)

				_snake_case			:      Union[str, Any]	    = joblib.load(os.path.join(__lowerCAmelCase						,      'args.joblib'			)			)
				_snake_case			:      int	    = joblib.load(os.path.join(__lowerCAmelCase						,      'data_collator.joblib'			)			)

				with open(os.path.join(__lowerCAmelCase						,      'training_state.json'			)						,      'r'			) as f:
								_snake_case			:      Dict	    = json.load(__lowerCAmelCase			)
				_snake_case			:      Tuple	    = training_state["""step"""]

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

def    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase			):
				_snake_case			:      Optional[int]	    = num_train_steps - warmup_steps
				_snake_case			:      Tuple	    = optax.linear_schedule(init_value=__lowerCAmelCase						,      end_value=__lowerCAmelCase						,      transition_steps=__lowerCAmelCase			)
				_snake_case			:      Optional[Any]	    = optax.linear_schedule(init_value=__lowerCAmelCase						,      end_value=1E-7						,      transition_steps=__lowerCAmelCase			)
				_snake_case			:      str	    = optax.join_schedules(schedules=[warmup_fn, decay_fn]						,      boundaries=[warmup_steps]			)
				return lr

def    A__(		__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase			):



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

				_snake_case			:      int	    = scheduler_fn(__lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase						,      __lowerCAmelCase			)

				_snake_case			:      List[Any]	    = optax.adamw(learning_rate=__lowerCAmelCase						,      weight_decay=__lowerCAmelCase						,      mask=__lowerCAmelCase			)
				return tx, lr













import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin
class     UpperCamelCase__					(     lowerCamelCase__				, lowerCamelCase__				, lowerCamelCase__				):






		'''simple docstring'''






		@register_to_config
		def __init__(							self,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__ = False,   )       -> Any:






									"""simple docstring"""
									super().__init__()

									lowercase_	:   str  				=     nn.Embedding(snake_case__,   snake_case__  )

									lowercase_	:   Dict  				=     nn.Embedding(snake_case__,   snake_case__  )
									lowercase_	:   List[str]  				=     False

									lowercase_	:   Optional[int]  				=     nn.Dropout(p=snake_case__  )

									lowercase_	:   Any  				=     TaConfig(
									    vocab_size=snake_case__,   d_model=snake_case__,   num_heads=snake_case__,   d_kv=snake_case__,   d_ff=snake_case__,   dropout_rate=snake_case__,   feed_forward_proj=snake_case__,   is_decoder=snake_case__,   is_encoder_decoder=snake_case__,   )

									lowercase_	:   List[Any]  				=     nn.ModuleList()
									for lyr_num in range(snake_case__  ):
																lowercase_	:   List[str]  				=     TaBlock(snake_case__  )
																self.encoders.append(snake_case__  )

									lowercase_	:   Tuple  				=     TaLayerNorm(snake_case__  )
									lowercase_	:   Dict  				=     nn.Dropout(p=snake_case__  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__  )       -> Union[str, Any]:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     self.token_embedder(snake_case__  )

									lowercase_	:   Tuple  				=     encoder_input_tokens.shape[1]
									lowercase_	:   List[str]  				=     torch.arange(snake_case__,   device=encoder_input_tokens.device  )
									x += self.position_encoding(snake_case__  )

									lowercase_	:   Union[str, Any]  				=     self.dropout_pre(snake_case__  )

									# inverted the attention mask
									lowercase_	:   List[str]  				=     encoder_input_tokens.size()
									lowercase_	:   Union[str, Any]  				=     self.get_extended_attention_mask(snake_case__,   snake_case__  )

									for lyr in self.encoders:
																lowercase_	:   Optional[Any]  				=     lyr(snake_case__,   snake_case__  )[0]
									lowercase_	:   Dict  				=     self.layer_norm(snake_case__  )

									return self.dropout_post(snake_case__  ), encoder_inputs_mask