
import logging
import os
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


UpperCamelCase_   	=	{
    "debug": logging.DEBUG,
    "info": logging.INFO,
    "warning": logging.WARNING,
    "error": logging.ERROR,
    "critical": logging.CRITICAL,
}

UpperCamelCase_   	=	logging.WARNING






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	__lowerCAmelCase         =							os.getenv("DATASETS_VERBOSITY"			,							UpperCamelCase  )
	if env_level_str:
		if env_level_str in log_levels:
			return log_levels[env_level_str]
		else:
			logging.getLogger().warning(
			    F"Unknown option DATASETS_VERBOSITY={env_level_str}, "
			    F"has to be one of: { ', '.join(log_levels.keys()  ) }"  )
	return _default_log_level






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	return __name__.split("."  )[0]






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	return logging.getLogger(_get_library_name()  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	__lowerCAmelCase         =							_get_library_root_logger()
	library_root_logger.setLevel(_get_default_logging_level()  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	__lowerCAmelCase         =							_get_library_root_logger()
	library_root_logger.setLevel(logging.NOTSET  )






def 		_UpperCAmelCase   (	UpperCamelCase:     Optional[str] = None  ):






	"""simple docstring"""


	if name is None:
		__lowerCAmelCase         =							_get_library_name()
	return logging.getLogger(UpperCamelCase  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	return _get_library_root_logger().getEffectiveLevel()






def 		_UpperCAmelCase   (	UpperCamelCase:     int  ):






	"""simple docstring"""


	_get_library_root_logger().setLevel(UpperCamelCase  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	return set_verbosity(UpperCamelCase  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	return set_verbosity(UpperCamelCase  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	return set_verbosity(UpperCamelCase  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	return set_verbosity(UpperCamelCase  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	__lowerCAmelCase         =							False






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	__lowerCAmelCase         =							True


# Configure the library root logger at the module level (singleton-like)
_configure_library_root_logger()

class       a		:







							def __init__(						self  :  List[str]      ,				*snake_case__  :  Optional[Any]      ,				**snake_case__  :  str ):  # pylint: disable=unused-argument
								"""simple docstring"""







								__lowerCAmelCase         =							args[0] if args else None







							def __iter__(						self  :  List[str] ):
								"""simple docstring"""







								return iter(self._iterator )







							def __getattr__(						self  :  Dict      ,				snake_case__  :  List[str] ):
								"""simple docstring"""







								def empty_fn(*snake_case__  :  str      ,				**snake_case__  :  Optional[Any] ):  # pylint: disable=unused-argument
									return

								return empty_fn







							def __enter__(						self  :  Optional[Any] ):
								"""simple docstring"""







								return self







							def __exit__(						self  :  List[Any]      ,				snake_case__  :  Optional[int]      ,				snake_case__  :  List[Any]      ,				snake_case__  :  Tuple ):
								"""simple docstring"""







								return
UpperCamelCase_   	=	True

class       a		:







							def __call__(						self  :  List[str]      ,				*snake_case__  :  List[str]      ,				snake_case__  :  Dict=False      ,				**snake_case__  :  int ):
								"""simple docstring"""







								if _tqdm_active and not disable:
									return tqdm_lib.tqdm(*snake_case__      ,				**snake_case__ )
								else:
									return EmptyTqdm(*snake_case__      ,				**snake_case__ )







							def       UpperCAmelCase__					(						self  :  List[str]      ,				*snake_case__  :  Optional[Any]      ,				**snake_case__  :  List[str] ):
								"""simple docstring"""







								__lowerCAmelCase         =							None
								if _tqdm_active:
									return tqdm_lib.tqdm.set_lock(*snake_case__      ,				**snake_case__ )







							def       UpperCAmelCase__					(						self  :  List[Any] ):
								"""simple docstring"""







								if _tqdm_active:
									return tqdm_lib.tqdm.get_lock()
UpperCamelCase_   	=	_tqdm_cls()






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	global _tqdm_active
	return bool(_tqdm_active  )






def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	global _tqdm_active
	__lowerCAmelCase         =							True







def 		_UpperCAmelCase   (	):






	"""simple docstring"""


	global _tqdm_active
	__lowerCAmelCase         =							False








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		(   __UpperCAmelCase					):
							lowercase_		:						BigBirdConfig
							lowercase_		:						jnp.dtype   							=  jnp.floataa
							lowercase_		:						bool   							=  True







							def       UpperCAmelCase__					(						self  :  Optional[int] ):
								"""simple docstring"""







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







							def __call__(						self  :  Optional[Any]      ,				*snake_case__  :  List[str]      ,				**snake_case__  :  str ):
								"""simple docstring"""







								__lowerCAmelCase         =							super().__call__(*snake_case__      ,				**snake_case__ )
								__lowerCAmelCase         =							self.cls(outputs[2] )
								return outputs[:2] + (cls_out,)

class       a		(   __UpperCAmelCase					):
							lowercase_		:						List[str]   							=  FlaxBigBirdForNaturalQuestionsModule






def 		_UpperCAmelCase   (	UpperCamelCase:     Optional[Any]			,							UpperCamelCase:     List[str]			,							UpperCamelCase:     Optional[Any]			,							UpperCamelCase:     List[Any]			,							UpperCamelCase:     Optional[Any]			,							UpperCamelCase:     int  ):






	"""simple docstring"""


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

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

@dataclass
class       a		:
							lowercase_		:						str   							=  "google/bigbird-roberta-base"
							lowercase_		:						int   							=  3_000
							lowercase_		:						int   							=  10_500

							lowercase_		:						int   							=  128
							lowercase_		:						int   							=  3

							lowercase_		:						int   							=  1
							lowercase_		:						int   							=  5

							# tx_args
							lowercase_		:						float   							=  3e-5
							lowercase_		:						float   							=  0.0
							lowercase_		:						int   							=  20_000
							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       UpperCAmelCase__					(						self  :  Optional[int] ):
								"""simple docstring"""







								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       a		:
							lowercase_		:						int
							lowercase_		:						int   							=  4_096  # no dynamic padding on TPUs







							def __call__(						self  :  List[Any]      ,				snake_case__  :  Union[str, Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							self.collate_fn(snake_case__ )
								__lowerCAmelCase         =							jax.tree_util.tree_map(snake_case__      ,				snake_case__ )
								return batch







							def       UpperCAmelCase__					(						self  :  List[Any]      ,				snake_case__  :  Dict ):
								"""simple docstring"""







								__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       UpperCAmelCase__					(						self  :  Optional[int]      ,				snake_case__  :  list ):
								"""simple docstring"""







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







							def       UpperCAmelCase__					(						self  :  Union[str, Any]      ,				snake_case__  :  list ):
								"""simple docstring"""







								__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   (	UpperCamelCase:     Tuple			,							UpperCamelCase:     Dict			,							UpperCamelCase:     Optional[Any]=None  ):






	"""simple docstring"""


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






@partial(jax.pmap			,							axis_name="batch"  )
def 		_UpperCAmelCase   (	UpperCamelCase:     Optional[int]			,							UpperCamelCase:     Union[str, Any]			,							**UpperCamelCase:     List[str]  ):






	"""simple docstring"""


	def loss_fn(UpperCamelCase:     Dict  ):
		__lowerCAmelCase         =							model_inputs.pop("start_labels"  )
		__lowerCAmelCase         =							model_inputs.pop("end_labels"  )
		__lowerCAmelCase         =							model_inputs.pop("pooled_labels"  )

		__lowerCAmelCase         =							state.apply_fn(**UpperCamelCase			,							params=UpperCamelCase			,							dropout_rng=UpperCamelCase			,							train=UpperCamelCase  )
		__lowerCAmelCase   ,						__lowerCAmelCase   ,						__lowerCAmelCase         =							outputs

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

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

	__lowerCAmelCase         =							state.apply_gradients(grads=UpperCamelCase  )
	return state, metrics, new_drp_rng






@partial(jax.pmap			,							axis_name="batch"  )
def 		_UpperCAmelCase   (	UpperCamelCase:     Optional[Any]			,							**UpperCamelCase:     List[str]  ):






	"""simple docstring"""


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

	__lowerCAmelCase         =							state.apply_fn(**UpperCamelCase			,							params=state.params			,							train=UpperCamelCase  )
	__lowerCAmelCase   ,						__lowerCAmelCase   ,						__lowerCAmelCase         =							outputs

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

class       a		(   train_state.TrainState					):
							lowercase_		:						Callable   							=  struct.field(pytree_node=__UpperCAmelCase					)

@dataclass
class       a		:
							lowercase_		:						Args
							lowercase_		:						Callable
							lowercase_		:						Callable
							lowercase_		:						Callable
							lowercase_		:						Callable
							lowercase_		:						wandb
							lowercase_		:						Callable   							=  None







							def       UpperCAmelCase__					(						self  :  Tuple      ,				snake_case__  :  int      ,				snake_case__  :  List[Any]      ,				snake_case__  :  Any      ,				snake_case__  :  str=None ):
								"""simple docstring"""







								__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       UpperCAmelCase__					(						self  :  str      ,				snake_case__  :  int      ,				snake_case__  :  Union[str, Any]      ,				snake_case__  :  Any ):
								"""simple docstring"""







								__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       UpperCAmelCase__					(						self  :  List[Any]      ,				snake_case__  :  Any      ,				snake_case__  :  Optional[Any] ):
								"""simple docstring"""







								__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       UpperCAmelCase__					(						self  :  Union[str, Any]      ,				snake_case__  :  Optional[int]      ,				snake_case__  :  Any ):
								"""simple docstring"""







								__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   (	UpperCamelCase:     str			,							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:
		__lowerCAmelCase         =							from_bytes(state.params			,							f.read()  )

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

	__lowerCAmelCase         =							joblib.load(os.path.join(UpperCamelCase			,							"args.joblib"  )  )
	__lowerCAmelCase         =							joblib.load(os.path.join(UpperCamelCase			,							"data_collator.joblib"  )  )

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

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






def 		_UpperCAmelCase   (	UpperCamelCase:     Any			,							UpperCamelCase:     Dict			,							UpperCamelCase:     Tuple			,							UpperCamelCase:     Dict  ):






	"""simple docstring"""


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







def 		_UpperCAmelCase   (	UpperCamelCase:     Union[str, Any]			,							UpperCamelCase:     str			,							UpperCamelCase:     List[Any]			,							UpperCamelCase:     Optional[int]			,							UpperCamelCase:     Union[str, Any]  ):






	"""simple docstring"""


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

	__lowerCAmelCase         =							scheduler_fn(UpperCamelCase			,							UpperCamelCase			,							UpperCamelCase			,							UpperCamelCase  )

	__lowerCAmelCase         =							optax.adamw(learning_rate=UpperCamelCase			,							weight_decay=UpperCamelCase			,							mask=UpperCamelCase  )
	return tx, lr






