


'''simple docstring'''

import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding





class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ):
     UpperCAmelCase 					=				['''image_processor''', '''tokenizer''']
     UpperCAmelCase 					=				'''ViTImageProcessor'''
     UpperCAmelCase 					=				('''CLIPTokenizer''', '''CLIPTokenizerFast''')






     def __init__(			self     ,							__UpperCamelCase=None     ,							__UpperCamelCase=None     ,							**__UpperCamelCase		)       ->    str:
          _a						=		None
          if "feature_extractor" in kwargs:
               warnings.warn(
                   "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
                   " instead."     ,							__UpperCamelCase     ,							)
               _a						=		kwargs.pop("feature_extractor"		)

          _a						=		image_processor if image_processor is not None else feature_extractor
          if image_processor is None:
               raise ValueError("You need to specify an `image_processor`."		)
          if tokenizer is None:
               raise ValueError("You need to specify a `tokenizer`."		)

          super().__init__(__UpperCamelCase     ,							__UpperCamelCase		)






     def __call__(			self     ,							__UpperCamelCase=None     ,							__UpperCamelCase=None     ,							__UpperCamelCase=None     ,							__UpperCamelCase=None     ,							**__UpperCamelCase		)       ->    Optional[int]:
          if text is None and visual_prompt is None and images is None:
               raise ValueError("You have to specify either text, visual prompt or images."		)

          if text is not None and visual_prompt is not None:
               raise ValueError("You have to specify exactly one type of prompt. Either text or visual prompt."		)

          if text is not None:
               _a						=		self.tokenizer(__UpperCamelCase     ,							return_tensors=__UpperCamelCase     ,							**__UpperCamelCase		)

          if visual_prompt is not None:
               _a						=		self.image_processor(__UpperCamelCase     ,							return_tensors=__UpperCamelCase     ,							**__UpperCamelCase		)

          if images is not None:
               _a						=		self.image_processor(__UpperCamelCase     ,							return_tensors=__UpperCamelCase     ,							**__UpperCamelCase		)

          if visual_prompt is not None and images is not None:
               _a						=		{
                   "pixel_values": image_features.pixel_values,
                   "conditional_pixel_values": prompt_features.pixel_values,
               }
               return encoding
          elif text is not None and images is not None:
               _a						=		image_features.pixel_values
               return encoding
          elif text is not None:
               return encoding
          elif visual_prompt is not None:
               _a						=		{
                   "conditional_pixel_values": prompt_features.pixel_values,
               }
               return encoding
          else:
               return BatchEncoding(data=dict(**__UpperCamelCase		)     ,							tensor_type=__UpperCamelCase		)






     def  a_					(			self     ,							*__UpperCamelCase     ,							**__UpperCamelCase		)       ->    str:
          return self.tokenizer.batch_decode(*__UpperCamelCase     ,							**__UpperCamelCase		)






     def  a_					(			self     ,							*__UpperCamelCase     ,							**__UpperCamelCase		)       ->    Tuple:
          return self.tokenizer.decode(*__UpperCamelCase     ,							**__UpperCamelCase		)






     @property
     def  a_					(			self		)       ->    str:
          warnings.warn(
              "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."     ,							__UpperCamelCase     ,							)
          return self.image_processor_class






     @property
     def  a_					(			self		)       ->    List[str]:
          warnings.warn(
              "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."     ,							__UpperCamelCase     ,							)
          return self.image_processor





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

