




"""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 lowerCAmelCase__			(  lowercase		):






			'''simple docstring'''



			lowerCamelCase__   =    42
			lowerCamelCase__   =    jnp.floataa
			lowerCamelCase__   =    True
			def       A_					(						self						):
							super().setup()
							_lowerCamelCase  :    str					=					nn.Dense(5						,   dtype=self.dtype						)






			def __call__(						self						,   *lowercase						,   **lowercase						):
							_lowerCamelCase  :    str					=					super().__call__(*lowercase						,   **lowercase						)
							_lowerCamelCase  :    Optional[Any]					=					self.cls(outputs[2]						)
							return outputs[:2] + (cls_out,)






class lowerCAmelCase__			(  lowercase		):






			'''simple docstring'''



			lowerCamelCase__   =    FlaxBigBirdForNaturalQuestionsModule







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

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






@dataclass
class lowerCAmelCase__			:






			'''simple docstring'''



			lowerCamelCase__   =    "google/bigbird-roberta-base"
			lowerCamelCase__   =    30_00
			lowerCamelCase__   =    1_05_00

			lowerCamelCase__   =    1_28
			lowerCamelCase__   =    3

			lowerCamelCase__   =    1
			lowerCamelCase__   =    5

			# tx_args
			lowerCamelCase__   =    3e-5
			lowerCamelCase__   =    0.0
			lowerCamelCase__   =    2_00_00
			lowerCamelCase__   =    0.0095

			lowerCamelCase__   =    "bigbird-roberta-natural-questions"
			lowerCamelCase__   =    "training-expt"
			lowerCamelCase__   =    "data/nq-training.jsonl"
			lowerCamelCase__   =    "data/nq-validation.jsonl"
			def       A_					(						self						):
							os.makedirs(self.base_dir						,   exist_ok=lowercase						)
							_lowerCamelCase  :    Optional[Any]					=					os.path.join(self.base_dir						,   self.save_dir						)
							_lowerCamelCase  :    List[Any]					=					self.batch_size_per_device * jax.device_count()






@dataclass
class lowerCAmelCase__			:






			'''simple docstring'''



			lowerCamelCase__   =    42
			lowerCamelCase__   =    40_96  # no dynamic padding on TPUs
			def __call__(						self						,   lowercase						):
							_lowerCamelCase  :    Dict					=					self.collate_fn(lowercase						)
							_lowerCamelCase  :    Union[str, Any]					=					jax.tree_util.tree_map(lowercase						,   lowercase						)
							return batch
			def       A_					(						self						,   lowercase						):
							_lowerCamelCase,				_lowerCamelCase  :    List[Any]					=					self.fetch_inputs(features['input_ids']						)
							_lowerCamelCase  :    Dict					=					{
							    'input_ids': jnp.array(lowercase						,   dtype=jnp.intaa						),
							    'attention_mask': jnp.array(lowercase						,   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						,   lowercase						):
							_lowerCamelCase  :    List[str]					=					[self._fetch_inputs(lowercase						) for ids in input_ids]
							return zip(*lowercase						)






			def       A_					(						self						,   lowercase						):
							_lowerCamelCase  :    Optional[Any]					=					[1 for _ in range(len(lowercase						)						)]
							while len(lowercase						) < self.max_length:
											input_ids.append(self.pad_id						)
											attention_mask.append(0						)
							return input_ids, attention_mask







def   _snake_case    (       lowercase__							,					lowercase__							,					lowercase__=None	):
				if seed is not None:
								_lowerCamelCase  :    Union[str, Any]					=					dataset.shuffle(seed=lowercase__	)
				for i in range(len(lowercase__	) // batch_size	):
								_lowerCamelCase  :    Dict					=					dataset[i * batch_size : (i + 1) * batch_size]
								yield dict(lowercase__	)







@partial(jax.pmap							,					axis_name='batch'	)
def   _snake_case    (       lowercase__							,					lowercase__							,					**lowercase__	):
				def loss_fn(lowercase__	):
								_lowerCamelCase  :    Optional[Any]					=					model_inputs.pop('start_labels'	)
								_lowerCamelCase  :    List[Any]					=					model_inputs.pop('end_labels'	)
								_lowerCamelCase  :    Any					=					model_inputs.pop('pooled_labels'	)

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

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

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

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







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

				_lowerCamelCase  :    str					=					state.apply_fn(**lowercase__							,					params=state.params							,					train=lowercase__	)
				_lowerCamelCase,				_lowerCamelCase,				_lowerCamelCase  :    Any					=					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 lowerCAmelCase__			(  train_state.TrainState		):






			'''simple docstring'''



			lowerCamelCase__   =    struct.field(pytree_node=lowercase		)






@dataclass
class lowerCAmelCase__			:






			'''simple docstring'''



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

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

																			_lowerCamelCase  :    List[Any]					=					self.evaluate(lowercase						,   lowercase						)
																			_lowerCamelCase  :    List[Any]					=					{
																			    'step': state_step.item(),
																			    'eval_loss': eval_loss.item(),
																			    'tr_loss': tr_loss,
																			    'lr': lr.item(),
																			}
																			tqdm.write(str(lowercase						)						)
																			self.logger.log(lowercase						,   commit=lowercase						)

															if i % args.save_steps == 0:
																			self.save_checkpoint(args.save_dir + F'''-e{epoch}-s{i}'''						,   state=lowercase						)
			def       A_					(						self						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    List[str]					=					get_batched_dataset(lowercase						,   self.args.batch_size						)
							_lowerCamelCase  :    Tuple					=					len(lowercase						) // self.args.batch_size
							_lowerCamelCase  :    Union[str, Any]					=					jnp.array(0						,   dtype=jnp.floataa						)
							_lowerCamelCase  :    Tuple					=					0
							for batch in tqdm(lowercase						,   total=lowercase						,   desc='Evaluating ... '						):
											_lowerCamelCase  :    List[Any]					=					self.data_collator(lowercase						)
											_lowerCamelCase  :    Tuple					=					self.val_step_fn(lowercase						,   **lowercase						)
											running_loss += jax_utils.unreplicate(metrics['loss']						)
											i += 1
							return running_loss / i






			def       A_					(						self						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    Any					=					jax_utils.unreplicate(lowercase						)
							print(F'''SAVING CHECKPOINT IN {save_dir}'''						,   end=' ... '						)
							self.model_save_fn(lowercase						,   params=state.params						)
							with open(os.path.join(lowercase						,   'opt_state.msgpack'						)						,   'wb'						) as f:
											f.write(to_bytes(state.opt_state						)						)
							joblib.dump(self.args						,   os.path.join(lowercase						,   'args.joblib'						)						)
							joblib.dump(self.data_collator						,   os.path.join(lowercase						,   'data_collator.joblib'						)						)
							with open(os.path.join(lowercase						,   'training_state.json'						)						,   'w'						) as f:
											json.dump({'step': state.step.item()}						,   lowercase						)
							print('DONE'						)







def   _snake_case    (       lowercase__							,					lowercase__	):
				print(f'''RESTORING CHECKPOINT FROM {save_dir}'''							,					end=' ... '	)
				with open(os.path.join(lowercase__							,					'flax_model.msgpack'	)							,					'rb'	) as f:
								_lowerCamelCase  :    int					=					from_bytes(state.params							,					f.read()	)

				with open(os.path.join(lowercase__							,					'opt_state.msgpack'	)							,					'rb'	) as f:
								_lowerCamelCase  :    Tuple					=					from_bytes(state.opt_state							,					f.read()	)

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

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

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







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







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

				_lowerCamelCase  :    List[Any]					=					scheduler_fn(lowercase__							,					lowercase__							,					lowercase__							,					lowercase__	)

				_lowerCamelCase  :    Any					=					optax.adamw(learning_rate=lowercase__							,					weight_decay=lowercase__							,					mask=lowercase__	)
				return tx, lr





"""simple docstring"""






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


lowercase__  		=		logging.get_logger(__name__)

lowercase__  		=		{
    """google/realm-cc-news-pretrained-embedder""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-encoder""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-scorer""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json"""
    ),
    """google/realm-cc-news-pretrained-openqa""": (
        """https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json"""
    ),
    """google/realm-orqa-nq-openqa""": """https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json""",
    """google/realm-orqa-nq-reader""": """https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json""",
    """google/realm-orqa-wq-openqa""": """https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json""",
    """google/realm-orqa-wq-reader""": """https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json""",
    # See all REALM models at https://huggingface.co/models?filter=realm
}






class lowerCAmelCase__			(  lowercase		):






			'''simple docstring'''



			lowerCamelCase__   =    """realm"""
			def __init__(						self						,   lowercase=30522						,   lowercase=768						,   lowercase=128						,   lowercase=12						,   lowercase=12						,   lowercase=8						,   lowercase=3072						,   lowercase="gelu_new"						,   lowercase=0.1						,   lowercase=0.1						,   lowercase=512						,   lowercase=2						,   lowercase=0.02						,   lowercase=1E-12						,   lowercase=256						,   lowercase=10						,   lowercase=1E-3						,   lowercase=5						,   lowercase=320						,   lowercase=13353718						,   lowercase=5000						,   lowercase=1						,   lowercase=0						,   lowercase=2						,   **lowercase						,   ):
							super().__init__(pad_token_id=lowercase						,   bos_token_id=lowercase						,   eos_token_id=lowercase						,   **lowercase						)

							# Common config
							_lowerCamelCase  :    str					=					vocab_size
							_lowerCamelCase  :    Dict					=					max_position_embeddings
							_lowerCamelCase  :    int					=					hidden_size
							_lowerCamelCase  :    Optional[Any]					=					retriever_proj_size
							_lowerCamelCase  :    Dict					=					num_hidden_layers
							_lowerCamelCase  :    Any					=					num_attention_heads
							_lowerCamelCase  :    int					=					num_candidates
							_lowerCamelCase  :    List[Any]					=					intermediate_size
							_lowerCamelCase  :    int					=					hidden_act
							_lowerCamelCase  :    Union[str, Any]					=					hidden_dropout_prob
							_lowerCamelCase  :    Dict					=					attention_probs_dropout_prob
							_lowerCamelCase  :    Union[str, Any]					=					initializer_range
							_lowerCamelCase  :    List[Any]					=					type_vocab_size
							_lowerCamelCase  :    int					=					layer_norm_eps

							# Reader config
							_lowerCamelCase  :    Tuple					=					span_hidden_size
							_lowerCamelCase  :    int					=					max_span_width
							_lowerCamelCase  :    Tuple					=					reader_layer_norm_eps
							_lowerCamelCase  :    Union[str, Any]					=					reader_beam_size
							_lowerCamelCase  :    Union[str, Any]					=					reader_seq_len

							# Retrieval config
							_lowerCamelCase  :    Optional[Any]					=					num_block_records
							_lowerCamelCase  :    str					=					searcher_beam_size