





'''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 			UpperCAmelCase					(      snake_case_       ):
      _lowercase:      BigBirdConfig
      _lowercase:      jnp.dtype											=       jnp.floataa
      _lowercase:      bool											=       True




      def     lowercase__   (      self      :			List[str]  )    ->				int:
            super().setup()
            _lowerCAmelCase												=							nn.Dense(5      ,      dtype=self.dtype  )




      def __call__(      self      :			List[Any]      ,      *__snake_case      :			int      ,      **__snake_case      :			int  )    ->				List[str]:
            _lowerCAmelCase												=							super().__call__(*__snake_case      ,      **__snake_case  )
            _lowerCAmelCase												=							self.cls(outputs[2]  )
            return outputs[:2] + (cls_out,)




class 			UpperCAmelCase					(      snake_case_       ):
      _lowercase:      Tuple											=       FlaxBigBirdForNaturalQuestionsModule




def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


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

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




@dataclass
class 			UpperCAmelCase					:
      _lowercase:      str											=       "google/bigbird-roberta-base"
      _lowercase:      int											=       3000
      _lowercase:      int											=       10500

      _lowercase:      int											=       128
      _lowercase:      int											=       3

      _lowercase:      int											=       1
      _lowercase:      int											=       5

      # tx_args
      _lowercase:      float											=       3E-5
      _lowercase:      float											=       0.0
      _lowercase:      int											=       20000
      _lowercase:      float											=       0.0095

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




      def     lowercase__   (      self      :			List[str]  )    ->				Optional[Any]:
            os.makedirs(self.base_dir      ,      exist_ok=__snake_case  )
            _lowerCAmelCase												=							os.path.join(self.base_dir      ,      self.save_dir  )
            _lowerCAmelCase												=							self.batch_size_per_device * jax.device_count()




@dataclass
class 			UpperCAmelCase					:
      _lowercase:      int
      _lowercase:      int											=       4096  # no dynamic padding on TPUs




      def __call__(      self      :			Optional[Any]      ,      __snake_case      :			int  )    ->				str:
            _lowerCAmelCase												=							self.collate_fn(__snake_case  )
            _lowerCAmelCase												=							jax.tree_util.tree_map(__snake_case      ,      __snake_case  )
            return batch




      def     lowercase__   (      self      :			Optional[Any]      ,      __snake_case      :			Any  )    ->				Tuple:
            _lowerCAmelCase					,							_lowerCAmelCase												=							self.fetch_inputs(features["""input_ids"""]  )
            _lowerCAmelCase												=							{
                """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     lowercase__   (      self      :			Optional[int]      ,      __snake_case      :			list  )    ->				List[Any]:
            _lowerCAmelCase												=							[self._fetch_inputs(__snake_case  ) for ids in input_ids]
            return zip(*__snake_case  )




      def     lowercase__   (      self      :			str      ,      __snake_case      :			list  )    ->				List[str]:
            _lowerCAmelCase												=							[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        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase=None  ):





      """simple docstring"""


      if seed is not None:
            _lowerCAmelCase												=							dataset.shuffle(seed=lowerCAmelCase  )
      for i in range(len(lowerCAmelCase  ) // batch_size  ):
            _lowerCAmelCase												=							dataset[i * batch_size : (i + 1) * batch_size]
            yield dict(lowerCAmelCase  )
@partial(jax.pmap							,     axis_name="""batch"""  )
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     **lowerCAmelCase  ):





      """simple docstring"""


      def loss_fn(lowerCAmelCase  ):
            _lowerCAmelCase												=							model_inputs.pop("""start_labels"""  )
            _lowerCAmelCase												=							model_inputs.pop("""end_labels"""  )
            _lowerCAmelCase												=							model_inputs.pop("""pooled_labels"""  )

            _lowerCAmelCase												=							state.apply_fn(**lowerCAmelCase							,     params=lowerCAmelCase							,     dropout_rng=lowerCAmelCase							,     train=lowerCAmelCase  )
            _lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase												=							outputs

            return state.loss_fn(
                lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     )

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

      _lowerCAmelCase												=							state.apply_gradients(grads=lowerCAmelCase  )
      return state, metrics, new_drp_rng
@partial(jax.pmap							,     axis_name="""batch"""  )
def        UpperCamelCase__				(		lowerCAmelCase							,     **lowerCAmelCase  ):





      """simple docstring"""


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

      _lowerCAmelCase												=							state.apply_fn(**lowerCAmelCase							,     params=state.params							,     train=lowerCAmelCase  )
      _lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase												=							outputs

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




class 			UpperCAmelCase					(      train_state.TrainState       ):
      _lowercase:      Callable											=       struct.field(pytree_node=snake_case_       )




@dataclass
class 			UpperCAmelCase					:
      _lowercase:      Args
      _lowercase:      Callable
      _lowercase:      Callable
      _lowercase:      Callable
      _lowercase:      Callable
      _lowercase:      wandb
      _lowercase:      Callable											=       None




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




      def     lowercase__   (      self      :			List[str]      ,      __snake_case      :			str      ,      __snake_case      :			Tuple      ,      __snake_case      :			Optional[Any]  )    ->				Any:
            _lowerCAmelCase												=							self.args
            _lowerCAmelCase												=							len(__snake_case  ) // args.batch_size

            _lowerCAmelCase												=							jax.random.PRNGKey(0  )
            _lowerCAmelCase												=							jax.random.split(__snake_case      ,      jax.device_count()  )
            for epoch in range(args.max_epochs  ):
                  _lowerCAmelCase												=							jnp.array(0      ,      dtype=jnp.floataa  )
                  _lowerCAmelCase												=							get_batched_dataset(__snake_case      ,      args.batch_size      ,      seed=__snake_case  )
                  _lowerCAmelCase												=							0
                  for batch in tqdm(__snake_case      ,      total=__snake_case      ,      desc=f"Running EPOCH-{epoch}"  ):
                        _lowerCAmelCase												=							self.data_collator(__snake_case  )
                        _lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase												=							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:
                              _lowerCAmelCase												=							jax_utils.unreplicate(state.step  )
                              _lowerCAmelCase												=							running_loss.item() / i
                              _lowerCAmelCase												=							self.scheduler_fn(state_step - 1  )

                              _lowerCAmelCase												=							self.evaluate(__snake_case      ,      __snake_case  )
                              _lowerCAmelCase												=							{
                                  """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     lowercase__   (      self      :			Union[str, Any]      ,      __snake_case      :			Optional[int]      ,      __snake_case      :			Dict  )    ->				List[str]:
            _lowerCAmelCase												=							get_batched_dataset(__snake_case      ,      self.args.batch_size  )
            _lowerCAmelCase												=							len(__snake_case  ) // self.args.batch_size
            _lowerCAmelCase												=							jnp.array(0      ,      dtype=jnp.floataa  )
            _lowerCAmelCase												=							0
            for batch in tqdm(__snake_case      ,      total=__snake_case      ,      desc="""Evaluating ... """  ):
                  _lowerCAmelCase												=							self.data_collator(__snake_case  )
                  _lowerCAmelCase												=							self.val_step_fn(__snake_case      ,      **__snake_case  )
                  running_loss += jax_utils.unreplicate(metrics["""loss"""]  )
                  i += 1
            return running_loss / i




      def     lowercase__   (      self      :			Dict      ,      __snake_case      :			Dict      ,      __snake_case      :			Dict  )    ->				Optional[Any]:
            _lowerCAmelCase												=							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        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


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

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

      _lowerCAmelCase												=							joblib.load(os.path.join(lowerCAmelCase							,     """args.joblib"""  )  )
      _lowerCAmelCase												=							joblib.load(os.path.join(lowerCAmelCase							,     """data_collator.joblib"""  )  )

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

      print("""DONE"""  )
      return params, opt_state, step, args, data_collator
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


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





def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


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

      _lowerCAmelCase												=							scheduler_fn(lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  )

      _lowerCAmelCase												=							optax.adamw(learning_rate=lowerCAmelCase							,     weight_decay=lowerCAmelCase							,     mask=lowerCAmelCase  )
      return tx, lr













'''simple docstring'''






import inspect
import os
import unittest

import torch

import accelerate
from accelerate import debug_launcher
from accelerate.test_utils import (
    execute_subprocess_async,
    require_cpu,
    require_huggingface_suite,
    require_multi_gpu,
    require_single_gpu,
)
from accelerate.utils import patch_environment




@require_huggingface_suite
class 			UpperCAmelCase					(      unittest.TestCase       ):




      def     lowercase__   (      self      :			int  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							inspect.getfile(accelerate.test_utils  )
            _lowerCAmelCase												=							os.path.sep.join(
                mod_file.split(os.path.sep  )[:-1] + ["""scripts""", """external_deps""", """test_metrics.py"""]  )

            from accelerate.test_utils.scripts.external_deps import test_metrics  # noqa: F401

            _lowerCAmelCase												=							test_metrics




      @require_cpu
      def     lowercase__   (      self      :			Optional[Any]  )    ->				Optional[Any]:
            debug_launcher(self.test_metrics.main      ,      num_processes=1  )




      @require_cpu
      def     lowercase__   (      self      :			Tuple  )    ->				Tuple:
            debug_launcher(self.test_metrics.main  )




      @require_single_gpu
      def     lowercase__   (      self      :			Union[str, Any]  )    ->				str:
            self.test_metrics.main()




      @require_multi_gpu
      def     lowercase__   (      self      :			str  )    ->				List[str]:
            print(f"Found {torch.cuda.device_count()} devices."  )
            _lowerCAmelCase												=							["""torchrun""", f"--nproc_per_node={torch.cuda.device_count()}", self.test_file_path]
            with patch_environment(omp_num_threads=1  ):
                  execute_subprocess_async(__snake_case      ,      env=os.environ.copy()  )






