




"""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__				(	__magic_name__   ):
			lowercase_    =	42
			lowercase_    =	jnp.floataa
			lowercase_    =	True


			def       a_     (      self      :    str):


						"""simple docstring"""




						super().setup()
						__UpperCAmelCase	: Optional[int]		=					nn.Dense(5      ,  dtype=self.dtype)


			def __call__(      self      :    str      ,  *UpperCamelCase_      :    int      ,  **UpperCamelCase_      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: str		=					super().__call__(*UpperCamelCase_      ,  **UpperCamelCase_)
						__UpperCAmelCase	: Dict		=					self.cls(outputs[2])
						return outputs[:2] + (cls_out,)






class      a__				(	__magic_name__   ):
			lowercase_    =	FlaxBigBirdForNaturalQuestionsModule






def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase							)			->			str:


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

			__UpperCAmelCase	: Optional[Any]		=					partial(UpperCamelCase				,       reduction=jnp.mean							)
			__UpperCAmelCase	: Tuple		=					cross_entropy(UpperCamelCase				,       UpperCamelCase							)
			__UpperCAmelCase	: Tuple		=					cross_entropy(UpperCamelCase				,       UpperCamelCase							)
			__UpperCAmelCase	: Optional[int]		=					cross_entropy(UpperCamelCase				,       UpperCamelCase							)
			return (start_loss + end_loss + pooled_loss) / 3



@dataclass
class      a__				:
			lowercase_    =	"google/bigbird-roberta-base"
			lowercase_    =	3_0_0_0
			lowercase_    =	1_0_5_0_0

			lowercase_    =	1_2_8
			lowercase_    =	3

			lowercase_    =	1
			lowercase_    =	5

			# tx_args
			lowercase_    =	3E-5
			lowercase_    =	0.0
			lowercase_    =	2_0_0_0_0
			lowercase_    =	0.00_95

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


			def       a_     (      self      :    int):


						"""simple docstring"""




						os.makedirs(self.base_dir      ,  exist_ok=UpperCamelCase_)
						__UpperCAmelCase	: Union[str, Any]		=					os.path.join(self.base_dir      ,  self.save_dir)
						__UpperCAmelCase	: List[Any]		=					self.batch_size_per_device * jax.device_count()






@dataclass
class      a__				:
			lowercase_    =	42
			lowercase_    =	4_0_9_6  # no dynamic padding on TPUs


			def __call__(      self      :    Optional[int]      ,  UpperCamelCase_      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					self.collate_fn(UpperCamelCase_)
						__UpperCAmelCase	: int		=					jax.tree_util.tree_map(UpperCamelCase_      ,  UpperCamelCase_)
						return batch


			def       a_     (      self      :    str      ,  UpperCamelCase_      :    Dict):


						"""simple docstring"""




						__UpperCAmelCase	,     __UpperCAmelCase	: Tuple		=					self.fetch_inputs(features["input_ids"])
						__UpperCAmelCase	: List[Any]		=					{
						    "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      :    Union[str, Any]      ,  UpperCamelCase_      :    list):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					[self._fetch_inputs(UpperCamelCase_) for ids in input_ids]
						return zip(*UpperCamelCase_)


			def       a_     (      self      :    Any      ,  UpperCamelCase_      :    list):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					[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    (			UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase=None							)			->			int:


			"""simple docstring"""
			if seed is not None:
						__UpperCAmelCase	: Tuple		=					dataset.shuffle(seed=UpperCamelCase							)
			for i in range(len(UpperCamelCase							) // batch_size							):
						__UpperCAmelCase	: List[str]		=					dataset[i * batch_size : (i + 1) * batch_size]
						yield dict(UpperCamelCase							)

@partial(jax.pmap				,       axis_name="batch"							)
def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase				,       **UpperCamelCase							)			->			Optional[int]:


			"""simple docstring"""
			def loss_fn(UpperCamelCase							):
						__UpperCAmelCase	: Optional[int]		=					model_inputs.pop("start_labels"							)
						__UpperCAmelCase	: Tuple		=					model_inputs.pop("end_labels"							)
						__UpperCAmelCase	: Optional[Any]		=					model_inputs.pop("pooled_labels"							)

						__UpperCAmelCase	: Union[str, Any]		=					state.apply_fn(**UpperCamelCase				,       params=UpperCamelCase				,       dropout_rng=UpperCamelCase				,       train=UpperCamelCase							)
						__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase	: Tuple		=					outputs

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

			__UpperCAmelCase	,     __UpperCAmelCase	: Optional[Any]		=					jax.random.split(UpperCamelCase							)
			__UpperCAmelCase	: Union[str, Any]		=					jax.value_and_grad(UpperCamelCase							)
			__UpperCAmelCase	,     __UpperCAmelCase	: Optional[Any]		=					grad_fn(state.params							)
			__UpperCAmelCase	: List[Any]		=					jax.lax.pmean({"loss": loss}				,       axis_name="batch"							)
			__UpperCAmelCase	: Any		=					jax.lax.pmean(UpperCamelCase				,       "batch"							)

			__UpperCAmelCase	: Tuple		=					state.apply_gradients(grads=UpperCamelCase							)
			return state, metrics, new_drp_rng

@partial(jax.pmap				,       axis_name="batch"							)
def 					_UpperCamelCase    (			UpperCamelCase				,       **UpperCamelCase							)			->			Any:


			"""simple docstring"""
			__UpperCAmelCase	: Optional[Any]		=					model_inputs.pop("start_labels"							)
			__UpperCAmelCase	: str		=					model_inputs.pop("end_labels"							)
			__UpperCAmelCase	: Union[str, Any]		=					model_inputs.pop("pooled_labels"							)

			__UpperCAmelCase	: Any		=					state.apply_fn(**UpperCamelCase				,       params=state.params				,       train=UpperCamelCase							)
			__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase	: str		=					outputs

			__UpperCAmelCase	: List[str]		=					state.loss_fn(UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase							)
			__UpperCAmelCase	: Union[str, Any]		=					jax.lax.pmean({"loss": loss}				,       axis_name="batch"							)
			return metrics



class      a__				(	train_state.TrainState   ):
			lowercase_    =	struct.field(pytree_node=__magic_name__   )






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


			def       a_     (      self      :    int      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    List[Any]=None):


						"""simple docstring"""




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


			def       a_     (      self      :    Union[str, Any]      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					self.args
						__UpperCAmelCase	: Any		=					len(UpperCamelCase_) // args.batch_size

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

															__UpperCAmelCase	: Any		=					self.evaluate(UpperCamelCase_      ,  UpperCamelCase_)
															__UpperCAmelCase	: List[Any]		=					{
															    "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      :    str      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					get_batched_dataset(UpperCamelCase_      ,  self.args.batch_size)
						__UpperCAmelCase	: List[Any]		=					len(UpperCamelCase_) // self.args.batch_size
						__UpperCAmelCase	: Tuple		=					jnp.array(0      ,  dtype=jnp.floataa)
						__UpperCAmelCase	: Optional[int]		=					0
						for batch in tqdm(UpperCamelCase_      ,  total=UpperCamelCase_      ,  desc="Evaluating ... "):
									__UpperCAmelCase	: List[Any]		=					self.data_collator(UpperCamelCase_)
									__UpperCAmelCase	: Union[str, Any]		=					self.val_step_fn(UpperCamelCase_      ,  **UpperCamelCase_)
									running_loss += jax_utils.unreplicate(metrics["loss"])
									i += 1
						return running_loss / i


			def       a_     (      self      :    List[Any]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					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    (			UpperCamelCase				,       UpperCamelCase							)			->			List[Any]:


			"""simple docstring"""
			print(f"RESTORING CHECKPOINT FROM {save_dir}"				,       end=" ... "							)
			with open(os.path.join(UpperCamelCase				,       "flax_model.msgpack"							)				,       "rb"							) as f:
						__UpperCAmelCase	: List[Any]		=					from_bytes(state.params				,       f.read()							)

			with open(os.path.join(UpperCamelCase				,       "opt_state.msgpack"							)				,       "rb"							) as f:
						__UpperCAmelCase	: Union[str, Any]		=					from_bytes(state.opt_state				,       f.read()							)

			__UpperCAmelCase	: Tuple		=					joblib.load(os.path.join(UpperCamelCase				,       "args.joblib"							)							)
			__UpperCAmelCase	: Union[str, Any]		=					joblib.load(os.path.join(UpperCamelCase				,       "data_collator.joblib"							)							)

			with open(os.path.join(UpperCamelCase				,       "training_state.json"							)				,       "r"							) as f:
						__UpperCAmelCase	: Optional[int]		=					json.load(UpperCamelCase							)
			__UpperCAmelCase	: List[str]		=					training_state["step"]

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

def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase							)			->			Optional[Any]:


			"""simple docstring"""
			__UpperCAmelCase	: Tuple		=					num_train_steps - warmup_steps
			__UpperCAmelCase	: int		=					optax.linear_schedule(init_value=UpperCamelCase				,       end_value=UpperCamelCase				,       transition_steps=UpperCamelCase							)
			__UpperCAmelCase	: Optional[int]		=					optax.linear_schedule(init_value=UpperCamelCase				,       end_value=1e-7				,       transition_steps=UpperCamelCase							)
			__UpperCAmelCase	: Optional[int]		=					optax.join_schedules(schedules=[warmup_fn, decay_fn]				,       boundaries=[warmup_steps]							)
			return lr





def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase							)			->			Any:


			"""simple docstring"""
			def weight_decay_mask(UpperCamelCase							):
						__UpperCAmelCase	: List[str]		=					traverse_util.flatten_dict(UpperCamelCase							)
						__UpperCAmelCase	: Union[str, Any]		=					{k: (v[-1] != "bias" and v[-2:] != ("LayerNorm", "scale")) for k, v in params.items()}
						return traverse_util.unflatten_dict(UpperCamelCase							)

			__UpperCAmelCase	: int		=					scheduler_fn(UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase				,       UpperCamelCase							)

			__UpperCAmelCase	: int		=					optax.adamw(learning_rate=UpperCamelCase				,       weight_decay=UpperCamelCase				,       mask=UpperCamelCase							)
			return tx, lr








"""simple docstring"""





import json
import os
import tempfile

from transformers.testing_utils import check_json_file_has_correct_format



class      a__				:
			lowercase_    =	None


			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					self.feature_extraction_class(**self.feat_extract_dict)
						__UpperCAmelCase	: Any		=					json.loads(feat_extract.to_json_string())
						for key, value in self.feat_extract_dict.items():
									self.assertEqual(obj[key]      ,  UpperCamelCase_)


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					self.feature_extraction_class(**self.feat_extract_dict)

						with tempfile.TemporaryDirectory() as tmpdirname:
									__UpperCAmelCase	: Dict		=					os.path.join(UpperCamelCase_      ,  "feat_extract.json")
									feat_extract_first.to_json_file(UpperCamelCase_)
									__UpperCAmelCase	: Optional[Any]		=					self.feature_extraction_class.from_json_file(UpperCamelCase_)

						self.assertEqual(feat_extract_second.to_dict()      ,  feat_extract_first.to_dict())


			def       a_     (      self      :    List[str]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					self.feature_extraction_class(**self.feat_extract_dict)

						with tempfile.TemporaryDirectory() as tmpdirname:
									__UpperCAmelCase	: Tuple		=					feat_extract_first.save_pretrained(UpperCamelCase_)[0]
									check_json_file_has_correct_format(UpperCamelCase_)
									__UpperCAmelCase	: List[Any]		=					self.feature_extraction_class.from_pretrained(UpperCamelCase_)

						self.assertEqual(feat_extract_second.to_dict()      ,  feat_extract_first.to_dict())


			def       a_     (      self      :    Union[str, Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					self.feature_extraction_class()
						self.assertIsNotNone(UpperCamelCase_)


