




'''simple docstring'''


from ....utils import logging


_snake_case				:				Union[str, Any]       							=		logging.get_logger(__name__)

class       A   (		_a   ):







							def __init__(  self	:  List[Any]							,      lowerCAmelCase_	:  List[Any]							,      lowerCAmelCase_	:  List[Any]=None							,      lowerCAmelCase_	:  str=20_48 )      ->							Dict:


														"""simple docstring"""


														_a     					=    config.__dict__
														_a     					=    modal_hidden_size
														if num_labels:
																					_a     					=    num_labels







'''simple docstring'''


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       A   (		_a   ):
							lowercase_          =    42
							lowercase_          =    jnp.floataa
							lowercase_          =    True







							def        __lowerCAmelCase      (  self	:  Optional[int] )      ->							Dict:


														"""simple docstring"""


														super().setup()
														_a     					=    nn.Dense(5							,      dtype=self.dtype )







							def __call__(  self	:  List[Any]							,      *lowerCAmelCase_	:  str							,      **lowerCAmelCase_	:  Dict )      ->							List[Any]:


														"""simple docstring"""


														_a     					=    super().__call__(*lowerCAmelCase_							,      **lowerCAmelCase_ )
														_a     					=    self.cls(outputs[2] )
														return outputs[:2] + (cls_out,)



class       A   (		_a   ):
							lowercase_          =    FlaxBigBirdForNaturalQuestionsModule



def 		snake_case_   (UpperCamelCase	:  Any			,						UpperCamelCase	:  str			,						UpperCamelCase	:  int			,						UpperCamelCase	:  Tuple			,						UpperCamelCase	:  Union[str, Any]			,						UpperCamelCase	:  int				):





							'''simple docstring'''


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

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

@dataclass
class       A   :
							lowercase_          =    "google/bigbird-roberta-base"
							lowercase_          =    3000
							lowercase_          =    1_0500

							lowercase_          =    128
							lowercase_          =    3

							lowercase_          =    1
							lowercase_          =    5

							# tx_args
							lowercase_          =    3e-5
							lowercase_          =    0.0
							lowercase_          =    2_0000
							lowercase_          =    0.0095

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







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


														"""simple docstring"""


														os.makedirs(self.base_dir							,      exist_ok=lowerCAmelCase_ )
														_a     					=    os.path.join(self.base_dir							,      self.save_dir )
														_a     					=    self.batch_size_per_device * jax.device_count()



@dataclass
class       A   :
							lowercase_          =    42
							lowercase_          =    4096  # no dynamic padding on TPUs







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


														"""simple docstring"""


														_a     					=    self.collate_fn(lowerCAmelCase_ )
														_a     					=    jax.tree_util.tree_map(lowerCAmelCase_							,      lowerCAmelCase_ )
														return batch







							def        __lowerCAmelCase      (  self	:  List[Any]							,      lowerCAmelCase_	:  List[Any] )      ->							int:


														"""simple docstring"""


														_a					,							_a     					=    self.fetch_inputs(features['''input_ids'''] )
														_a     					=    {
														    '''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        __lowerCAmelCase      (  self	:  Dict							,      lowerCAmelCase_	:  list )      ->							List[Any]:


														"""simple docstring"""


														_a     					=    [self._fetch_inputs(lowerCAmelCase_ ) for ids in input_ids]
														return zip(*lowerCAmelCase_ )







							def        __lowerCAmelCase      (  self	:  Dict							,      lowerCAmelCase_	:  list )      ->							str:


														"""simple docstring"""


														_a     					=    [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 		snake_case_   (UpperCamelCase	:  List[Any]			,						UpperCamelCase	:  List[Any]			,						UpperCamelCase	:  Dict=None				):





							'''simple docstring'''


							if seed is not None:
														_a     					=    dataset.shuffle(seed=UpperCamelCase				)
							for i in range(len(UpperCamelCase				) // batch_size				):
														_a     					=    dataset[i * batch_size : (i + 1) * batch_size]
														yield dict(UpperCamelCase				)
@partial(jax.pmap			,						axis_name='''batch'''				)
def 		snake_case_   (UpperCamelCase	:  Optional[int]			,						UpperCamelCase	:  Tuple			,						**UpperCamelCase	:  str				):





							'''simple docstring'''


							def loss_fn(UpperCamelCase	:  List[str]				):
														_a     					=    model_inputs.pop('''start_labels'''				)
														_a     					=    model_inputs.pop('''end_labels'''				)
														_a     					=    model_inputs.pop('''pooled_labels'''				)

														_a     					=    state.apply_fn(**UpperCamelCase			,						params=UpperCamelCase			,						dropout_rng=UpperCamelCase			,						train=UpperCamelCase				)
														_a					,							_a					,							_a     					=    outputs

														return state.loss_fn(
														    UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						)

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

							_a     					=    state.apply_gradients(grads=UpperCamelCase				)
							return state, metrics, new_drp_rng
@partial(jax.pmap			,						axis_name='''batch'''				)
def 		snake_case_   (UpperCamelCase	:  Any			,						**UpperCamelCase	:  Optional[int]				):





							'''simple docstring'''


							_a     					=    model_inputs.pop('''start_labels'''				)
							_a     					=    model_inputs.pop('''end_labels'''				)
							_a     					=    model_inputs.pop('''pooled_labels'''				)

							_a     					=    state.apply_fn(**UpperCamelCase			,						params=state.params			,						train=UpperCamelCase				)
							_a					,							_a					,							_a     					=    outputs

							_a     					=    state.loss_fn(UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase				)
							_a     					=    jax.lax.pmean({'''loss''': loss}			,						axis_name='''batch'''				)
							return metrics

class       A   (		train_state.TrainState   ):
							lowercase_          =    struct.field(pytree_node=_a   )



@dataclass
class       A   :
							lowercase_          =    42
							lowercase_          =    42
							lowercase_          =    42
							lowercase_          =    42
							lowercase_          =    42
							lowercase_          =    42
							lowercase_          =    None







							def        __lowerCAmelCase      (  self	:  Tuple							,      lowerCAmelCase_	:  Optional[Any]							,      lowerCAmelCase_	:  Optional[Any]							,      lowerCAmelCase_	:  Optional[Any]							,      lowerCAmelCase_	:  Optional[int]=None )      ->							List[str]:


														"""simple docstring"""


														_a     					=    model.params
														_a     					=    TrainState.create(
														    apply_fn=model.__call__							,      params=lowerCAmelCase_							,      tx=lowerCAmelCase_							,      loss_fn=lowerCAmelCase_							,      )
														if ckpt_dir is not None:
																					_a					,							_a					,							_a					,							_a					,							_a     					=    restore_checkpoint(lowerCAmelCase_							,      lowerCAmelCase_ )
																					_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(**lowerCAmelCase_ )
																					_a     					=    train_state.TrainState(
																					    step=lowerCAmelCase_							,      apply_fn=model.__call__							,      params=lowerCAmelCase_							,      tx=lowerCAmelCase_							,      opt_state=lowerCAmelCase_							,      )
																					_a     					=    args
																					_a     					=    data_collator
																					_a     					=    lr
																					_a     					=    params
														_a     					=    jax_utils.replicate(lowerCAmelCase_ )
														return state







							def        __lowerCAmelCase      (  self	:  str							,      lowerCAmelCase_	:  Any							,      lowerCAmelCase_	:  Union[str, Any]							,      lowerCAmelCase_	:  Optional[Any] )      ->							int:


														"""simple docstring"""


														_a     					=    self.args
														_a     					=    len(lowerCAmelCase_ ) // args.batch_size

														_a     					=    jax.random.PRNGKey(0 )
														_a     					=    jax.random.split(lowerCAmelCase_							,      jax.device_count() )
														for epoch in range(args.max_epochs ):
																					_a     					=    jnp.array(0							,      dtype=jnp.floataa )
																					_a     					=    get_batched_dataset(lowerCAmelCase_							,      args.batch_size							,      seed=lowerCAmelCase_ )
																					_a     					=    0
																					for batch in tqdm(lowerCAmelCase_							,      total=lowerCAmelCase_							,      desc=F'Running EPOCH-{epoch}' ):
																												_a     					=    self.data_collator(lowerCAmelCase_ )
																												_a					,							_a					,							_a     					=    self.train_step_fn(lowerCAmelCase_							,      lowerCAmelCase_							,      **lowerCAmelCase_ )
																												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(lowerCAmelCase_							,      lowerCAmelCase_ )
																																			_a     					=    {
																																			    '''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        __lowerCAmelCase      (  self	:  List[str]							,      lowerCAmelCase_	:  List[str]							,      lowerCAmelCase_	:  str )      ->							List[str]:


														"""simple docstring"""


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







							def        __lowerCAmelCase      (  self	:  int							,      lowerCAmelCase_	:  List[Any]							,      lowerCAmelCase_	:  str )      ->							int:


														"""simple docstring"""


														_a     					=    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 		snake_case_   (UpperCamelCase	:  int			,						UpperCamelCase	:  Dict				):





							'''simple docstring'''


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

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

							_a     					=    joblib.load(os.path.join(UpperCamelCase			,						'''args.joblib'''				)				)
							_a     					=    joblib.load(os.path.join(UpperCamelCase			,						'''data_collator.joblib'''				)				)

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

							print('''DONE'''				)
							return params, opt_state, step, args, data_collator
def 		snake_case_   (UpperCamelCase	:  Dict			,						UpperCamelCase	:  Optional[int]			,						UpperCamelCase	:  Dict			,						UpperCamelCase	:  List[Any]				):





							'''simple docstring'''


							_a     					=    num_train_steps - warmup_steps
							_a     					=    optax.linear_schedule(init_value=UpperCamelCase			,						end_value=UpperCamelCase			,						transition_steps=UpperCamelCase				)
							_a     					=    optax.linear_schedule(init_value=UpperCamelCase			,						end_value=1e-7			,						transition_steps=UpperCamelCase				)
							_a     					=    optax.join_schedules(schedules=[warmup_fn, decay_fn]			,						boundaries=[warmup_steps]				)
							return lr

def 		snake_case_   (UpperCamelCase	:  Tuple			,						UpperCamelCase	:  Optional[Any]			,						UpperCamelCase	:  str			,						UpperCamelCase	:  Union[str, Any]			,						UpperCamelCase	:  int				):





							'''simple docstring'''


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

							_a     					=    scheduler_fn(UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase				)

							_a     					=    optax.adamw(learning_rate=UpperCamelCase			,						weight_decay=UpperCamelCase			,						mask=UpperCamelCase				)
							return tx, lr

