


'''simple docstring'''

import argparse
import os

import numpy as np
import tensorflow as tf
import torch

from transformers import BertModel






def __UpperCamelCase   (			__lowerCamelCase		: BertModel      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: str				)      ->     List[str]:







     '''simple docstring'''






     _a						=		("dense.weight", "attention.self.query", "attention.self.key", "attention.self.value")

     _a						=		(
         ("layer.", "layer_"),
         ("word_embeddings.weight", "word_embeddings"),
         ("position_embeddings.weight", "position_embeddings"),
         ("token_type_embeddings.weight", "token_type_embeddings"),
         (".", "/"),
         ("LayerNorm/weight", "LayerNorm/gamma"),
         ("LayerNorm/bias", "LayerNorm/beta"),
         ("weight", "kernel"),
     )

     if not os.path.isdir(__lowerCamelCase				):
          os.makedirs(__lowerCamelCase				)

     _a						=		model.state_dict()

     def to_tf_var_name(__lowerCamelCase		: str				):
          for patt, repl in iter(__lowerCamelCase				):
               _a						=		name.replace(__lowerCamelCase      ,  __lowerCamelCase				)
          return F"bert/{name}"

     def create_tf_var(__lowerCamelCase		: np.ndarray      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: tf.Session				):
          _a						=		tf.dtypes.as_dtype(tensor.dtype				)
          _a						=		tf.get_variable(dtype=__lowerCamelCase      ,  shape=tensor.shape      ,  name=__lowerCamelCase      ,  initializer=tf.zeros_initializer()				)
          session.run(tf.variables_initializer([tf_var]				)				)
          session.run(__lowerCamelCase				)
          return tf_var

     tf.reset_default_graph()
     with tf.Session() as session:
          for var_name in state_dict:
               _a						=		to_tf_var_name(__lowerCamelCase				)
               _a						=		state_dict[var_name].numpy()
               if any(x in var_name for x in tensors_to_transpose				):
                    _a						=		torch_tensor.T
               _a						=		create_tf_var(tensor=__lowerCamelCase      ,  name=__lowerCamelCase      ,  session=__lowerCamelCase				)
               tf.keras.backend.set_value(__lowerCamelCase      ,  __lowerCamelCase				)
               _a						=		session.run(__lowerCamelCase				)
               print(F"Successfully created {tf_name}: {np.allclose(__lowerCamelCase      ,  __lowerCamelCase				)}"				)

          _a						=		tf.train.Saver(tf.trainable_variables()				)
          saver.save(__lowerCamelCase      ,  os.path.join(__lowerCamelCase      ,  model_name.replace("-"      ,  "_"				) + ".ckpt"				)				)






def __UpperCamelCase   (			__lowerCamelCase		: str=None				)      ->     Optional[int]:







     '''simple docstring'''






     _a						=		argparse.ArgumentParser()
     parser.add_argument("--model_name"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="model name e.g. bert-base-uncased"				)
     parser.add_argument(
         "--cache_dir"      ,  type=__lowerCamelCase      ,  default=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="Directory containing pytorch model"				)
     parser.add_argument("--pytorch_model_path"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="/path/to/<pytorch-model-name>.bin"				)
     parser.add_argument("--tf_cache_dir"      ,  type=__lowerCamelCase      ,  required=__lowerCamelCase      ,  help="Directory in which to save tensorflow model"				)
     _a						=		parser.parse_args(__lowerCamelCase				)

     _a						=		BertModel.from_pretrained(
         pretrained_model_name_or_path=args.model_name      ,  state_dict=torch.load(args.pytorch_model_path				)      ,  cache_dir=args.cache_dir      ,  )

     convert_pytorch_checkpoint_to_tf(model=__lowerCamelCase      ,  ckpt_dir=args.tf_cache_dir      ,  model_name=args.model_name				)


if __name__ == "__main__":
   main()





'''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 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ):
     UpperCAmelCase 					=				42
     UpperCAmelCase 					=				jnp.floataa
     UpperCAmelCase 					=				True






     def  a_					(			self		)       ->    int:
          super().setup()
          _a						=		nn.Dense(5     ,							dtype=self.dtype		)






     def __call__(			self     ,							*__UpperCamelCase     ,							**__UpperCamelCase		)       ->    Optional[Any]:
          _a						=		super().__call__(*__UpperCamelCase     ,							**__UpperCamelCase		)
          _a						=		self.cls(outputs[2]		)
          return outputs[:2] + (cls_out,)







class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ):
     UpperCAmelCase 					=				FlaxBigBirdForNaturalQuestionsModule







def __UpperCamelCase   (			__lowerCamelCase		: str      ,  __lowerCamelCase		: Any      ,  __lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: Dict				)      ->     List[Any]:







     '''simple docstring'''






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

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





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

     UpperCAmelCase 					=				128
     UpperCAmelCase 					=				3

     UpperCAmelCase 					=				1
     UpperCAmelCase 					=				5

     # tx_args
     UpperCAmelCase 					=				3E-5
     UpperCAmelCase 					=				0.0
     UpperCAmelCase 					=				20_000
     UpperCAmelCase 					=				0.00_95

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






     def  a_					(			self		)       ->    Tuple:
          os.makedirs(self.base_dir     ,							exist_ok=__UpperCamelCase		)
          _a						=		os.path.join(self.base_dir     ,							self.save_dir		)
          _a						=		self.batch_size_per_device * jax.device_count()







@dataclass
class 						__SCREAMING_SNAKE_CASE					:
     UpperCAmelCase 					=				42
     UpperCAmelCase 					=				4_096  # no dynamic padding on TPUs






     def __call__(			self     ,							__UpperCamelCase		)       ->    Dict:
          _a						=		self.collate_fn(__UpperCamelCase		)
          _a						=		jax.tree_util.tree_map(__UpperCamelCase     ,							__UpperCamelCase		)
          return batch






     def  a_					(			self     ,							__UpperCamelCase		)       ->    Tuple:
          _a			,			_a						=		self.fetch_inputs(features["input_ids"]		)
          _a						=		{
              "input_ids": jnp.array(__UpperCamelCase     ,							dtype=jnp.intaa		),
              "attention_mask": jnp.array(__UpperCamelCase     ,							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     ,							__UpperCamelCase		)       ->    List[str]:
          _a						=		[self._fetch_inputs(__UpperCamelCase		) for ids in input_ids]
          return zip(*__UpperCamelCase		)






     def  a_					(			self     ,							__UpperCamelCase		)       ->    List[str]:
          _a						=		[1 for _ in range(len(__UpperCamelCase		)		)]
          while len(__UpperCamelCase		) < self.max_length:
               input_ids.append(self.pad_id		)
               attention_mask.append(0		)
          return input_ids, attention_mask







def __UpperCamelCase   (			__lowerCamelCase		: List[str]      ,  __lowerCamelCase		: Optional[Any]      ,  __lowerCamelCase		: List[Any]=None				)      ->     List[Any]:







     '''simple docstring'''






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






@partial(jax.pmap      ,  axis_name="batch"				)
def __UpperCamelCase   (			__lowerCamelCase		: Union[str, Any]      ,  __lowerCamelCase		: int      ,  **__lowerCamelCase		: Optional[int]				)      ->     Any:







     '''simple docstring'''






     def loss_fn(__lowerCamelCase		: str				):
          _a						=		model_inputs.pop("start_labels"				)
          _a						=		model_inputs.pop("end_labels"				)
          _a						=		model_inputs.pop("pooled_labels"				)

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

          return state.loss_fn(
              __lowerCamelCase      ,  __lowerCamelCase      ,  __lowerCamelCase      ,  __lowerCamelCase      ,  __lowerCamelCase      ,  __lowerCamelCase      ,  )

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

     _a						=		state.apply_gradients(grads=__lowerCamelCase				)
     return state, metrics, new_drp_rng






@partial(jax.pmap      ,  axis_name="batch"				)
def __UpperCamelCase   (			__lowerCamelCase		: Any      ,  **__lowerCamelCase		: Tuple				)      ->     List[Any]:







     '''simple docstring'''






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

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

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





class 						__SCREAMING_SNAKE_CASE					(      train_state.TrainState ):
     UpperCAmelCase 					=				struct.field(pytree_node=lowerCamelCase__ )







@dataclass
class 						__SCREAMING_SNAKE_CASE					:
     UpperCAmelCase 					=				42
     UpperCAmelCase 					=				42
     UpperCAmelCase 					=				42
     UpperCAmelCase 					=				42
     UpperCAmelCase 					=				42
     UpperCAmelCase 					=				42
     UpperCAmelCase 					=				None






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






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    str:
          _a						=		self.args
          _a						=		len(__UpperCamelCase		) // args.batch_size

          _a						=		jax.random.PRNGKey(0		)
          _a						=		jax.random.split(__UpperCamelCase     ,							jax.device_count()		)
          for epoch in range(args.max_epochs		):
               _a						=		jnp.array(0     ,							dtype=jnp.floataa		)
               _a						=		get_batched_dataset(__UpperCamelCase     ,							args.batch_size     ,							seed=__UpperCamelCase		)
               _a						=		0
               for batch in tqdm(__UpperCamelCase     ,							total=__UpperCamelCase     ,							desc=f"Running EPOCH-{epoch}"		):
                    _a						=		self.data_collator(__UpperCamelCase		)
                    _a			,			_a			,			_a						=		self.train_step_fn(__UpperCamelCase     ,							__UpperCamelCase     ,							**__UpperCamelCase		)
                    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(__UpperCamelCase     ,							__UpperCamelCase		)
                         _a						=		{
                             "step": state_step.item(),
                             "eval_loss": eval_loss.item(),
                             "tr_loss": tr_loss,
                             "lr": lr.item(),
                         }
                         tqdm.write(str(__UpperCamelCase		)		)
                         self.logger.log(__UpperCamelCase     ,							commit=__UpperCamelCase		)

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






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    List[Any]:
          _a						=		get_batched_dataset(__UpperCamelCase     ,							self.args.batch_size		)
          _a						=		len(__UpperCamelCase		) // self.args.batch_size
          _a						=		jnp.array(0     ,							dtype=jnp.floataa		)
          _a						=		0
          for batch in tqdm(__UpperCamelCase     ,							total=__UpperCamelCase     ,							desc="Evaluating ... "		):
               _a						=		self.data_collator(__UpperCamelCase		)
               _a						=		self.val_step_fn(__UpperCamelCase     ,							**__UpperCamelCase		)
               running_loss += jax_utils.unreplicate(metrics["loss"]		)
               i += 1
          return running_loss / i






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Dict:
          _a						=		jax_utils.unreplicate(__UpperCamelCase		)
          print(f"SAVING CHECKPOINT IN {save_dir}"     ,							end=" ... "		)
          self.model_save_fn(__UpperCamelCase     ,							params=state.params		)
          with open(os.path.join(__UpperCamelCase     ,							"opt_state.msgpack"		)     ,							"wb"		) as f:
               f.write(to_bytes(state.opt_state		)		)
          joblib.dump(self.args     ,							os.path.join(__UpperCamelCase     ,							"args.joblib"		)		)
          joblib.dump(self.data_collator     ,							os.path.join(__UpperCamelCase     ,							"data_collator.joblib"		)		)
          with open(os.path.join(__UpperCamelCase     ,							"training_state.json"		)     ,							"w"		) as f:
               json.dump({"step": state.step.item()}     ,							__UpperCamelCase		)
          print("DONE"		)







def __UpperCamelCase   (			__lowerCamelCase		: Union[str, Any]      ,  __lowerCamelCase		: Dict				)      ->     Optional[int]:







     '''simple docstring'''






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

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

     _a						=		joblib.load(os.path.join(__lowerCamelCase      ,  "args.joblib"				)				)
     _a						=		joblib.load(os.path.join(__lowerCamelCase      ,  "data_collator.joblib"				)				)

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

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






def __UpperCamelCase   (			__lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: List[str]      ,  __lowerCamelCase		: List[str]				)      ->     List[str]:







     '''simple docstring'''






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






def __UpperCamelCase   (			__lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: str      ,  __lowerCamelCase		: List[Any]      ,  __lowerCamelCase		: List[str]      ,  __lowerCamelCase		: Tuple				)      ->     Any:







     '''simple docstring'''






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

     _a						=		scheduler_fn(__lowerCamelCase      ,  __lowerCamelCase      ,  __lowerCamelCase      ,  __lowerCamelCase				)

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

