



'''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  snake_case__						(    SCREAMING_SNAKE_CASE_				):
 A__				= 42
 A__				= jnp.floataa
 A__				= True
 def A_   (   self						:							Optional[int]  )      ->   List[str]:







       '''simple docstring'''
       super().setup()
       __snake_case :   Optional[Any]													=    nn.Dense(5      ,     dtype=self.dtype  )





 def __call__(   self						:							int      ,     *__a						:							Optional[int]      ,     **__a						:							int  )      ->   str:







       '''simple docstring'''
       __snake_case :   Optional[int]													=    super().__call__(*__a      ,     **__a  )
       __snake_case :   Union[str, Any]													=    self.cls(outputs[2]  )
       return outputs[:2] + (cls_out,)



class  snake_case__						(    SCREAMING_SNAKE_CASE_				):
 A__				= FlaxBigBirdForNaturalQuestionsModule





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

      __snake_case :   int													=    partial(_UpperCAmelCase ,reduction=jnp.mean				)
      __snake_case :   List[str]													=    cross_entropy(_UpperCAmelCase ,_UpperCAmelCase				)
      __snake_case :   Optional[int]													=    cross_entropy(_UpperCAmelCase ,_UpperCAmelCase				)
      __snake_case :   Optional[int]													=    cross_entropy(_UpperCAmelCase ,_UpperCAmelCase				)
      return (start_loss + end_loss + pooled_loss) / 3



@dataclass
class  snake_case__						:
 A__				= "google/bigbird-roberta-base"
 A__				= 3_000
 A__				= 10_500

 A__				= 128
 A__				= 3

 A__				= 1
 A__				= 5

 # tx_args
 A__				= 3E-5
 A__				= 0.0
 A__				= 20_000
 A__				= 0.0_0_9_5

 A__				= "bigbird-roberta-natural-questions"
 A__				= "training-expt"
 A__				= "data/nq-training.jsonl"
 A__				= "data/nq-validation.jsonl"
 def A_   (   self						:							Tuple  )      ->   Optional[int]:







       '''simple docstring'''
       os.makedirs(self.base_dir      ,     exist_ok=__a  )
       __snake_case :   Optional[Any]													=    os.path.join(self.base_dir      ,     self.save_dir  )
       __snake_case :   str													=    self.batch_size_per_device * jax.device_count()



@dataclass
class  snake_case__						:
 A__				= 42
 A__				= 4_096  # no dynamic padding on TPUs
 def __call__(   self						:							Optional[int]      ,     __a						:							Optional[int]  )      ->   Any:







       '''simple docstring'''
       __snake_case :   int													=    self.collate_fn(__a  )
       __snake_case :   Dict													=    jax.tree_util.tree_map(__a      ,     __a  )
       return batch
 def A_   (   self						:							str      ,     __a						:							Dict  )      ->   Any:







       '''simple docstring'''
       __snake_case							,     __snake_case :   Optional[int]													=    self.fetch_inputs(features['input_ids']  )
       __snake_case :   Optional[int]													=    {
           'input_ids': jnp.array(__a      ,     dtype=jnp.intaa  ),
           'attention_mask': jnp.array(__a      ,     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						:							Tuple      ,     __a						:							list  )      ->   int:







       '''simple docstring'''
       __snake_case :   Optional[Any]													=    [self._fetch_inputs(__a  ) for ids in input_ids]
       return zip(*__a  )





 def A_   (   self						:							Optional[int]      ,     __a						:							list  )      ->   Optional[int]:







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





def   a_   (		_UpperCAmelCase    :      Tuple ,_UpperCAmelCase    :      Dict ,_UpperCAmelCase    :      Optional[Any]=None				)    ->  List[str]:
      if seed is not None:
            __snake_case :   List[str]													=    dataset.shuffle(seed=_UpperCAmelCase				)
      for i in range(len(_UpperCAmelCase				) // batch_size				):
            __snake_case :   Optional[Any]													=    dataset[i * batch_size : (i + 1) * batch_size]
            yield dict(_UpperCAmelCase				)





@partial(jax.pmap ,axis_name='batch'				)
def   a_   (		_UpperCAmelCase    :      Optional[int] ,_UpperCAmelCase    :      Optional[Any] ,**_UpperCAmelCase    :      List[str]				)    ->  Optional[Any]:
      def loss_fn(_UpperCAmelCase    :      Optional[int]				):
            __snake_case :   Optional[int]													=    model_inputs.pop('start_labels'				)
            __snake_case :   List[str]													=    model_inputs.pop('end_labels'				)
            __snake_case :   int													=    model_inputs.pop('pooled_labels'				)

            __snake_case :   Any													=    state.apply_fn(**_UpperCAmelCase ,params=_UpperCAmelCase ,dropout_rng=_UpperCAmelCase ,train=_UpperCAmelCase				)
            __snake_case							,     __snake_case							,     __snake_case :   Union[str, Any]													=    outputs

            return state.loss_fn(
                _UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,)

      __snake_case							,     __snake_case :   Any													=    jax.random.split(_UpperCAmelCase				)
      __snake_case :   str													=    jax.value_and_grad(_UpperCAmelCase				)
      __snake_case							,     __snake_case :   Tuple													=    grad_fn(state.params				)
      __snake_case :   str													=    jax.lax.pmean({'loss': loss} ,axis_name='batch'				)
      __snake_case :   Union[str, Any]													=    jax.lax.pmean(_UpperCAmelCase ,'batch'				)

      __snake_case :   List[str]													=    state.apply_gradients(grads=_UpperCAmelCase				)
      return state, metrics, new_drp_rng





@partial(jax.pmap ,axis_name='batch'				)
def   a_   (		_UpperCAmelCase    :      List[Any] ,**_UpperCAmelCase    :      Optional[Any]				)    ->  Dict:
      __snake_case :   Tuple													=    model_inputs.pop('start_labels'				)
      __snake_case :   Union[str, Any]													=    model_inputs.pop('end_labels'				)
      __snake_case :   Dict													=    model_inputs.pop('pooled_labels'				)

      __snake_case :   Optional[int]													=    state.apply_fn(**_UpperCAmelCase ,params=state.params ,train=_UpperCAmelCase				)
      __snake_case							,     __snake_case							,     __snake_case :   List[str]													=    outputs

      __snake_case :   int													=    state.loss_fn(_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase				)
      __snake_case :   Union[str, Any]													=    jax.lax.pmean({'loss': loss} ,axis_name='batch'				)
      return metrics



class  snake_case__						(    train_state.TrainState				):
 A__				= struct.field(pytree_node=SCREAMING_SNAKE_CASE_				)



@dataclass
class  snake_case__						:
 A__				= 42
 A__				= 42
 A__				= 42
 A__				= 42
 A__				= 42
 A__				= 42
 A__				= None
 def A_   (   self						:							str      ,     __a						:							Optional[Any]      ,     __a						:							Optional[Any]      ,     __a						:							int      ,     __a						:							str=None  )      ->   List[str]:







       '''simple docstring'''
       __snake_case :   List[Any]													=    model.params
       __snake_case :   Any													=    TrainState.create(
           apply_fn=model.__call__      ,     params=__a      ,     tx=__a      ,     loss_fn=__a      ,     )
       if ckpt_dir is not None:
             __snake_case							,     __snake_case							,     __snake_case							,     __snake_case							,     __snake_case :   Union[str, Any]													=    restore_checkpoint(__a      ,     __a  )
             __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							,     __snake_case :   Tuple													=    build_tx(**__a  )
             __snake_case :   Tuple													=    train_state.TrainState(
                 step=__a      ,     apply_fn=model.__call__      ,     params=__a      ,     tx=__a      ,     opt_state=__a      ,     )
             __snake_case :   Dict													=    args
             __snake_case :   int													=    data_collator
             __snake_case :   Tuple													=    lr
             __snake_case :   List[Any]													=    params
       __snake_case :   Optional[int]													=    jax_utils.replicate(__a  )
       return state
 def A_   (   self						:							Tuple      ,     __a						:							Optional[int]      ,     __a						:							int      ,     __a						:							Union[str, Any]  )      ->   Union[str, Any]:







       '''simple docstring'''
       __snake_case :   Union[str, Any]													=    self.args
       __snake_case :   List[Any]													=    len(__a  ) // args.batch_size

       __snake_case :   str													=    jax.random.PRNGKey(0  )
       __snake_case :   str													=    jax.random.split(__a      ,     jax.device_count()  )
       for epoch in range(args.max_epochs  ):
             __snake_case :   Tuple													=    jnp.array(0      ,     dtype=jnp.floataa  )
             __snake_case :   Optional[Any]													=    get_batched_dataset(__a      ,     args.batch_size      ,     seed=__a  )
             __snake_case :   str													=    0
             for batch in tqdm(__a      ,     total=__a      ,     desc=f'''Running EPOCH-{epoch}'''  ):
                   __snake_case :   int													=    self.data_collator(__a  )
                   __snake_case							,     __snake_case							,     __snake_case :   Optional[Any]													=    self.train_step_fn(__a      ,     __a      ,     **__a  )
                   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[str]													=    running_loss.item() / i
                         __snake_case :   Tuple													=    self.scheduler_fn(state_step - 1  )

                         __snake_case :   Optional[Any]													=    self.evaluate(__a      ,     __a  )
                         __snake_case :   Any													=    {
                             'step': state_step.item(),
                             'eval_loss': eval_loss.item(),
                             'tr_loss': tr_loss,
                             'lr': lr.item(),
                         }
                         tqdm.write(str(__a  )  )
                         self.logger.log(__a      ,     commit=__a  )

                   if i % args.save_steps == 0:
                         self.save_checkpoint(args.save_dir + f'''-e{epoch}-s{i}'''      ,     state=__a  )
 def A_   (   self						:							int      ,     __a						:							Union[str, Any]      ,     __a						:							List[str]  )      ->   Dict:







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





 def A_   (   self						:							Optional[Any]      ,     __a						:							Optional[Any]      ,     __a						:							Dict  )      ->   Dict:







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





def   a_   (		_UpperCAmelCase    :      Any ,_UpperCAmelCase    :      List[str]				)    ->  Any:
      print(f'''RESTORING CHECKPOINT FROM {save_dir}''' ,end=' ... '				)
      with open(os.path.join(_UpperCAmelCase ,'flax_model.msgpack'				) ,'rb'				) as f:
            __snake_case :   List[Any]													=    from_bytes(state.params ,f.read()				)

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

      __snake_case :   Union[str, Any]													=    joblib.load(os.path.join(_UpperCAmelCase ,'args.joblib'				)				)
      __snake_case :   List[Any]													=    joblib.load(os.path.join(_UpperCAmelCase ,'data_collator.joblib'				)				)

      with open(os.path.join(_UpperCAmelCase ,'training_state.json'				) ,'r'				) as f:
            __snake_case :   Any													=    json.load(_UpperCAmelCase				)
      __snake_case :   Optional[int]													=    training_state['step']

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





def   a_   (		_UpperCAmelCase    :      Union[str, Any] ,_UpperCAmelCase    :      Optional[Any] ,_UpperCAmelCase    :      str ,_UpperCAmelCase    :      Tuple				)    ->  List[str]:
      __snake_case :   Union[str, Any]													=    num_train_steps - warmup_steps
      __snake_case :   Optional[Any]													=    optax.linear_schedule(init_value=_UpperCAmelCase ,end_value=_UpperCAmelCase ,transition_steps=_UpperCAmelCase				)
      __snake_case :   List[Any]													=    optax.linear_schedule(init_value=_UpperCAmelCase ,end_value=1E-7 ,transition_steps=_UpperCAmelCase				)
      __snake_case :   Optional[int]													=    optax.join_schedules(schedules=[warmup_fn, decay_fn] ,boundaries=[warmup_steps]				)
      return lr





def   a_   (		_UpperCAmelCase    :      Any ,_UpperCAmelCase    :      Optional[int] ,_UpperCAmelCase    :      List[Any] ,_UpperCAmelCase    :      Optional[Any] ,_UpperCAmelCase    :      Union[str, Any]				)    ->  List[Any]:
      def weight_decay_mask(_UpperCAmelCase    :      Optional[Any]				):
            __snake_case :   Union[str, Any]													=    traverse_util.flatten_dict(_UpperCAmelCase				)
            __snake_case :   Optional[Any]													=    {k: (v[-1] != 'bias' and v[-2:] != ('LayerNorm', 'scale')) for k, v in params.items()}
            return traverse_util.unflatten_dict(_UpperCAmelCase				)

      __snake_case :   Optional[Any]													=    scheduler_fn(_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase ,_UpperCAmelCase				)

      __snake_case :   str													=    optax.adamw(learning_rate=_UpperCAmelCase ,weight_decay=_UpperCAmelCase ,mask=_UpperCAmelCase				)
      return tx, lr








'''simple docstring'''





import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING


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



class  snake_case__						(    SCREAMING_SNAKE_CASE_				):
 A__				= '''upernet'''
 def __init__(   self						:							Optional[int]      ,     __a						:							Dict=None      ,     __a						:							List[str]=512      ,     __a						:							int=0.0_2      ,     __a						:							Optional[Any]=[1, 2, 3, 6]      ,     __a						:							Union[str, Any]=True      ,     __a						:							Any=0.4      ,     __a						:							int=384      ,     __a						:							int=256      ,     __a						:							str=1      ,     __a						:							List[Any]=False      ,     __a						:							Dict=255      ,     **__a						:							Dict      ,     )      ->   Optional[int]:







       '''simple docstring'''
       super().__init__(**__a  )

       if backbone_config is None:
             logger.info('`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.'  )
             __snake_case :   Optional[int]													=    CONFIG_MAPPING['resnet'](out_features=['stage1', 'stage2', 'stage3', 'stage4']  )
       elif isinstance(__a      ,     __a  ):
             __snake_case :   List[Any]													=    backbone_config.get('model_type'  )
             __snake_case :   List[str]													=    CONFIG_MAPPING[backbone_model_type]
             __snake_case :   int													=    config_class.from_dict(__a  )

       __snake_case :   Any													=    backbone_config
       __snake_case :   Tuple													=    hidden_size
       __snake_case :   str													=    initializer_range
       __snake_case :   Dict													=    pool_scales
       __snake_case :   Optional[Any]													=    use_auxiliary_head
       __snake_case :   Optional[Any]													=    auxiliary_loss_weight
       __snake_case :   Optional[int]													=    auxiliary_in_channels
       __snake_case :   Dict													=    auxiliary_channels
       __snake_case :   Any													=    auxiliary_num_convs
       __snake_case :   Optional[Any]													=    auxiliary_concat_input
       __snake_case :   List[str]													=    loss_ignore_index





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







       '''simple docstring'''
       __snake_case :   Union[str, Any]													=    copy.deepcopy(self.__dict__  )
       __snake_case :   str													=    self.backbone_config.to_dict()
       __snake_case :   Optional[int]													=    self.__class__.model_type
       return output



