

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      (_UpperCAmelCase ):
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  jnp.floataa
 _UpperCamelCase			  =  True
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       super().setup()
       __lowerCAmelCase					:						int					=					nn.Dense(5      ,							dtype=self.dtype					)





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



       '''simple docstring'''





       __lowerCAmelCase					:						Union[str, Any]					=					super().__call__(*A_      ,							**A_					)
       __lowerCAmelCase					:						List[str]					=					self.cls(outputs[2]					)
       return outputs[:2] + (cls_out,)







class   __lowercase      (_UpperCAmelCase ):
 _UpperCamelCase			  =  FlaxBigBirdForNaturalQuestionsModule







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

      __lowerCAmelCase					:						str					=					partial(lowercase__   ,	reduction=jnp.mean						)
      __lowerCAmelCase					:						List[str]					=					cross_entropy(lowercase__   ,	lowercase__						)
      __lowerCAmelCase					:						str					=					cross_entropy(lowercase__   ,	lowercase__						)
      __lowerCAmelCase					:						Dict					=					cross_entropy(lowercase__   ,	lowercase__						)
      return (start_loss + end_loss + pooled_loss) / 3







@dataclass
class   __lowercase      :
 _UpperCamelCase			  =  "google/bigbird-roberta-base"
 _UpperCamelCase			  =  3000
 _UpperCamelCase			  =  10500

 _UpperCamelCase			  =  128
 _UpperCamelCase			  =  3

 _UpperCamelCase			  =  1
 _UpperCamelCase			  =  5

 # tx_args
 _UpperCamelCase			  =  3E-5
 _UpperCamelCase			  =  0.0
 _UpperCamelCase			  =  20000
 _UpperCamelCase			  =  0.0095

 _UpperCamelCase			  =  "bigbird-roberta-natural-questions"
 _UpperCamelCase			  =  "training-expt"
 _UpperCamelCase			  =  "data/nq-training.jsonl"
 _UpperCamelCase			  =  "data/nq-validation.jsonl"
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





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







@dataclass
class   __lowercase      :
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  4096  # no dynamic padding on TPUs
 def __call__(							self      ,							A_					)					->Tuple:



       '''simple docstring'''





       __lowerCAmelCase					:						Union[str, Any]					=					self.collate_fn(A_					)
       __lowerCAmelCase					:						int					=					jax.tree_util.tree_map(A_      ,							A_					)
       return batch
 def 	UpperCamelCase__  (							self      ,							A_					)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase,							__lowerCAmelCase					:						Optional[Any]					=					self.fetch_inputs(features['''input_ids''']					)
       __lowerCAmelCase					:						Dict					=					{
           '''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 	UpperCamelCase__  (							self      ,							A_					)					->Union[str, Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						Any					=					[self._fetch_inputs(A_					) for ids in input_ids]
       return zip(*A_					)





 def 	UpperCamelCase__  (							self      ,							A_					)					->Dict:



       '''simple docstring'''





       __lowerCAmelCase					:						List[Any]					=					[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 		_lowercase (     lowercase__   ,	lowercase__   ,	lowercase__=None						):
      if seed is not None:
            __lowerCAmelCase					:						List[str]					=					dataset.shuffle(seed=lowercase__						)
      for i in range(len(lowercase__						) // batch_size						):
            __lowerCAmelCase					:						Tuple					=					dataset[i * batch_size : (i + 1) * batch_size]
            yield dict(lowercase__						)







@partial(jax.pmap   ,	axis_name='''batch'''						)
def 		_lowercase (     lowercase__   ,	lowercase__   ,	**lowercase__						):
      def loss_fn(lowercase__						):
            __lowerCAmelCase					:						Any					=					model_inputs.pop('''start_labels'''						)
            __lowerCAmelCase					:						Tuple					=					model_inputs.pop('''end_labels'''						)
            __lowerCAmelCase					:						str					=					model_inputs.pop('''pooled_labels'''						)

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

            return state.loss_fn(
                lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__   ,	)

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

      __lowerCAmelCase					:						Optional[int]					=					state.apply_gradients(grads=lowercase__						)
      return state, metrics, new_drp_rng







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

      __lowerCAmelCase					:						Dict					=					state.apply_fn(**lowercase__   ,	params=state.params   ,	train=lowercase__						)
      __lowerCAmelCase,							__lowerCAmelCase,							__lowerCAmelCase					:						Tuple					=					outputs

      __lowerCAmelCase					:						Optional[int]					=					state.loss_fn(lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__						)
      __lowerCAmelCase					:						List[str]					=					jax.lax.pmean({'''loss''': loss}   ,	axis_name='''batch'''						)
      return metrics







class   __lowercase      (train_state.TrainState ):
 _UpperCamelCase			  =  struct.field(pytree_node=_UpperCAmelCase )







@dataclass
class   __lowercase      :
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  None
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_=None					)					->Optional[int]:



       '''simple docstring'''





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



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					self.args
       __lowerCAmelCase					:						Dict					=					len(A_					) // args.batch_size

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

                         __lowerCAmelCase					:						Union[str, Any]					=					self.evaluate(A_      ,							A_					)
                         __lowerCAmelCase					:						int					=					{
                             '''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 	UpperCamelCase__  (							self      ,							A_      ,							A_					)					->Optional[Any]:



       '''simple docstring'''





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





 def 	UpperCamelCase__  (							self      ,							A_      ,							A_					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						Union[str, 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 		_lowercase (     lowercase__   ,	lowercase__						):
      print(f"""RESTORING CHECKPOINT FROM {save_dir}"""   ,	end=''' ... '''						)
      with open(os.path.join(lowercase__   ,	'''flax_model.msgpack'''						)   ,	'''rb'''						) as f:
            __lowerCAmelCase					:						str					=					from_bytes(state.params   ,	f.read()						)

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

      __lowerCAmelCase					:						Any					=					joblib.load(os.path.join(lowercase__   ,	'''args.joblib'''						)						)
      __lowerCAmelCase					:						List[Any]					=					joblib.load(os.path.join(lowercase__   ,	'''data_collator.joblib'''						)						)

      with open(os.path.join(lowercase__   ,	'''training_state.json'''						)   ,	'''r'''						) as f:
            __lowerCAmelCase					:						int					=					json.load(lowercase__						)
      __lowerCAmelCase					:						List[Any]					=					training_state['''step''']

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







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







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

      __lowerCAmelCase					:						Optional[int]					=					scheduler_fn(lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__						)

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





import unicodedata
from dataclasses import dataclass
from typing import Optional, Union

import numpy as np

from transformers.data.data_collator import DataCollatorMixin
from transformers.file_utils import PaddingStrategy
from transformers.tokenization_utils_base import PreTrainedTokenizerBase







def 		_lowercase (     lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__						):
      if isinstance(lowercase__   ,	lowercase__						):
            __lowerCAmelCase					:						Dict					=					np.full((len(lowercase__						), sequence_length, 2)   ,	lowercase__						)
      else:
            __lowerCAmelCase					:						Optional[int]					=					np.full((len(lowercase__						), sequence_length)   ,	lowercase__						)

      for i, tensor in enumerate(lowercase__						):
            if padding_side == "right":
                  if isinstance(lowercase__   ,	lowercase__						):
                        __lowerCAmelCase					:						Union[str, Any]					=					tensor[:sequence_length]
                  else:
                        __lowerCAmelCase					:						int					=					tensor[:sequence_length]
            else:
                  if isinstance(lowercase__   ,	lowercase__						):
                        __lowerCAmelCase					:						Union[str, Any]					=					tensor[:sequence_length]
                  else:
                        __lowerCAmelCase					:						Optional[Any]					=					tensor[:sequence_length]

      return out_tensor.tolist()







def 		_lowercase (     lowercase__						):
      __lowerCAmelCase					:						Union[str, Any]					=					ord(lowercase__						)
      if (cp >= 3_3 and cp <= 4_7) or (cp >= 5_8 and cp <= 6_4) or (cp >= 9_1 and cp <= 9_6) or (cp >= 1_2_3 and cp <= 1_2_6):
            return True
      __lowerCAmelCase					:						int					=					unicodedata.category(lowercase__						)
      if cat.startswith('''P'''						):
            return True
      return False







@dataclass
class   __lowercase      (_UpperCAmelCase ):
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  True
 _UpperCamelCase			  =  None
 _UpperCamelCase			  =  None
 _UpperCamelCase			  =  -100
 _UpperCamelCase			  =  "pt"
 def 	UpperCamelCase__  (							self      ,							A_					)					->Optional[int]:



       '''simple docstring'''





       import torch

       __lowerCAmelCase					:						List[str]					=					'''label''' if '''label''' in features[0].keys() else '''labels'''
       __lowerCAmelCase					:						Union[str, Any]					=					[feature[label_name] for feature in features] if label_name in features[0].keys() else None
       __lowerCAmelCase					:						List[Any]					=					self.tokenizer.pad(
           A_      ,							padding=self.padding      ,							max_length=self.max_length      ,							pad_to_multiple_of=self.pad_to_multiple_of      ,							return_tensors='''pt''' if labels is None else None      ,							)

       if labels is None:
             return batch

       __lowerCAmelCase					:						Dict					=					torch.tensor(batch['''entity_ids''']					).shape[1]
       __lowerCAmelCase					:						Optional[int]					=					self.tokenizer.padding_side
       if padding_side == "right":
             __lowerCAmelCase					:						Any					=					[
                 list(A_					) + [self.label_pad_token_id] * (sequence_length - len(A_					)) for label in labels
             ]
       else:
             __lowerCAmelCase					:						Optional[int]					=					[
                 [self.label_pad_token_id] * (sequence_length - len(A_					)) + list(A_					) for label in labels
             ]

       __lowerCAmelCase					:						Tuple					=					[feature['''ner_tags'''] for feature in features]
       __lowerCAmelCase					:						List[Any]					=					padding_tensor(A_      ,							-1      ,							A_      ,							A_					)
       __lowerCAmelCase					:						Optional[int]					=					[feature['''original_entity_spans'''] for feature in features]
       __lowerCAmelCase					:						Any					=					padding_tensor(A_      ,							(-1, -1)      ,							A_      ,							A_					)
       __lowerCAmelCase					:						Optional[Any]					=					{k: torch.tensor(A_      ,							dtype=torch.intaa					) for k, v in batch.items()}

       return batch


