


'''simple docstring'''

from typing import List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase__				     =		logging.get_logger(__name__)

lowercase__				     =		{
    "huggingface/autoformer-tourism-monthly": "https://huggingface.co/huggingface/autoformer-tourism-monthly/resolve/main/config.json",
}





class 						__SCREAMING_SNAKE_CASE					(      __SCREAMING_SNAKE_CASE ):
     UpperCAmelCase 					=				'''autoformer'''
     UpperCAmelCase 					=				{
         '''hidden_size''': '''d_model''',
         '''num_attention_heads''': '''encoder_attention_heads''',
         '''num_hidden_layers''': '''encoder_layers''',
     }






     def __init__(			self     ,							__UpperCamelCase = None     ,							__UpperCamelCase = None     ,							__UpperCamelCase = "student_t"     ,							__UpperCamelCase = "nll"     ,							__UpperCamelCase = 1     ,							__UpperCamelCase = [1, 2, 3, 4, 5, 6, 7]     ,							__UpperCamelCase = True     ,							__UpperCamelCase = 0     ,							__UpperCamelCase = 0     ,							__UpperCamelCase = 0     ,							__UpperCamelCase = 0     ,							__UpperCamelCase = None     ,							__UpperCamelCase = None     ,							__UpperCamelCase = 64     ,							__UpperCamelCase = 2     ,							__UpperCamelCase = 2     ,							__UpperCamelCase = 2     ,							__UpperCamelCase = 2     ,							__UpperCamelCase = 32     ,							__UpperCamelCase = 32     ,							__UpperCamelCase = "gelu"     ,							__UpperCamelCase = 0.1     ,							__UpperCamelCase = 0.1     ,							__UpperCamelCase = 0.1     ,							__UpperCamelCase = 0.1     ,							__UpperCamelCase = 0.1     ,							__UpperCamelCase = 100     ,							__UpperCamelCase = 0.02     ,							__UpperCamelCase = True     ,							__UpperCamelCase=True     ,							__UpperCamelCase = 10     ,							__UpperCamelCase = 25     ,							__UpperCamelCase = 3     ,							**__UpperCamelCase     ,							)       ->    List[Any]:
          # time series specific configuration
          _a						=		prediction_length
          _a						=		context_length if context_length is not None else prediction_length
          _a						=		distribution_output
          _a						=		loss
          _a						=		input_size
          _a						=		num_time_features
          _a						=		lags_sequence
          _a						=		scaling
          _a						=		num_dynamic_real_features
          _a						=		num_static_real_features
          _a						=		num_static_categorical_features
          if cardinality is not None and num_static_categorical_features > 0:
               if len(__snake_case		) != num_static_categorical_features:
                    raise ValueError(
                        "The cardinality should be a list of the same length as `num_static_categorical_features`"		)
               _a						=		cardinality
          else:
               _a						=		[0]
          if embedding_dimension is not None and num_static_categorical_features > 0:
               if len(__snake_case		) != num_static_categorical_features:
                    raise ValueError(
                        "The embedding dimension should be a list of the same length as `num_static_categorical_features`"		)
               _a						=		embedding_dimension
          else:
               _a						=		[min(50     ,							(cat + 1) // 2		) for cat in self.cardinality]
          _a						=		num_parallel_samples

          # Transformer architecture configuration
          _a						=		input_size * len(self.lags_sequence		) + self._number_of_features
          _a						=		d_model
          _a						=		encoder_attention_heads
          _a						=		decoder_attention_heads
          _a						=		encoder_ffn_dim
          _a						=		decoder_ffn_dim
          _a						=		encoder_layers
          _a						=		decoder_layers

          _a						=		dropout
          _a						=		attention_dropout
          _a						=		activation_dropout
          _a						=		encoder_layerdrop
          _a						=		decoder_layerdrop

          _a						=		activation_function
          _a						=		init_std

          _a						=		use_cache

          # Autoformer
          _a						=		label_length
          _a						=		moving_average
          _a						=		autocorrelation_factor

          super().__init__(is_encoder_decoder=__snake_case     ,							**__snake_case		)






     @property
     def  a_					(			self		)       ->    Tuple:
          return (
              sum(self.embedding_dimension		)
              + self.num_dynamic_real_features
              + self.num_time_features
              + self.num_static_real_features
              + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
          )





'''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

