


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     UpperCamelCase(  _a   ):
     snake_case_		:				BigBirdConfig
     snake_case_		:				jnp.dtype										=	jnp.floataa
     snake_case_		:				bool										=	True


     def 							SCREAMING_SNAKE_CASE_  (						self	:		str						)			->	Any:




           '''simple docstring'''

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






     def __call__(						self	:		Optional[int]     ,    *SCREAMING_SNAKE_CASE	:		Tuple     ,    **SCREAMING_SNAKE_CASE	:		List[Any]						)			->	str:




           '''simple docstring'''

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






class     UpperCamelCase(  _a   ):
     snake_case_		:				Optional[Any]										=	FlaxBigBirdForNaturalQuestionsModule





def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase )						->	List[Any]:


      '''simple docstring'''






      def cross_entropy(_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase=None ):
            __snake_case				=		logits.shape[-1]
            __snake_case				=		(labels[..., None] == jnp.arange(_lowerCAmelCase )[None]).astype("f4" )
            __snake_case				=		jax.nn.log_softmax(_lowerCAmelCase      ,	axis=-1 )
            __snake_case				=		-jnp.sum(labels * logits      ,	axis=-1 )
            if reduction is not None:
                  __snake_case				=		reduction(_lowerCAmelCase )
            return loss

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






@dataclass
class     UpperCamelCase:
     snake_case_		:				str										=	"google/bigbird-roberta-base"
     snake_case_		:				int										=	3000
     snake_case_		:				int										=	1_0500

     snake_case_		:				int										=	128
     snake_case_		:				int										=	3

     snake_case_		:				int										=	1
     snake_case_		:				int										=	5

     # tx_args
     snake_case_		:				float										=	3E-5
     snake_case_		:				float										=	0.0
     snake_case_		:				int										=	2_0000
     snake_case_		:				float										=	0.0095

     snake_case_		:				str										=	"bigbird-roberta-natural-questions"
     snake_case_		:				str										=	"training-expt"
     snake_case_		:				str										=	"data/nq-training.jsonl"
     snake_case_		:				str										=	"data/nq-validation.jsonl"


     def 							SCREAMING_SNAKE_CASE_  (						self	:		Union[str, Any]						)			->	Any:




           '''simple docstring'''

           os.makedirs(self.base_dir     ,    exist_ok=SCREAMING_SNAKE_CASE						)
           __snake_case				=		os.path.join(self.base_dir     ,    self.save_dir						)
           __snake_case				=		self.batch_size_per_device * jax.device_count()






@dataclass
class     UpperCamelCase:
     snake_case_		:				int
     snake_case_		:				int										=	4096  # no dynamic padding on TPUs


     def __call__(						self	:		int     ,    SCREAMING_SNAKE_CASE	:		Optional[Any]						)			->	Union[str, Any]:




           '''simple docstring'''

           __snake_case				=		self.collate_fn(SCREAMING_SNAKE_CASE						)
           __snake_case				=		jax.tree_util.tree_map(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE						)
           return batch


     def 							SCREAMING_SNAKE_CASE_  (						self	:		Optional[Any]     ,    SCREAMING_SNAKE_CASE	:		List[Any]						)			->	Optional[int]:




           '''simple docstring'''

           __snake_case		,					__snake_case				=		self.fetch_inputs(features["input_ids"]						)
           __snake_case				=		{
               "input_ids": jnp.array(SCREAMING_SNAKE_CASE     ,    dtype=jnp.intaa						),
               "attention_mask": jnp.array(SCREAMING_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 							SCREAMING_SNAKE_CASE_  (						self	:		Optional[int]     ,    SCREAMING_SNAKE_CASE	:		list						)			->	str:




           '''simple docstring'''

           __snake_case				=		[self._fetch_inputs(SCREAMING_SNAKE_CASE						) for ids in input_ids]
           return zip(*SCREAMING_SNAKE_CASE						)






     def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple     ,    SCREAMING_SNAKE_CASE	:		list						)			->	Optional[Any]:




           '''simple docstring'''

           __snake_case				=		[1 for _ in range(len(SCREAMING_SNAKE_CASE						)						)]
           while len(SCREAMING_SNAKE_CASE						) < self.max_length:
                 input_ids.append(self.pad_id						)
                 attention_mask.append(0						)
           return input_ids, attention_mask





def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase=None )						->	List[Any]:


      '''simple docstring'''






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





@partial(jax.pmap      ,	axis_name="batch" )
def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	_lowerCAmelCase      ,	**_lowerCAmelCase )						->	List[Any]:


      '''simple docstring'''






      def loss_fn(_lowerCAmelCase ):
            __snake_case				=		model_inputs.pop("start_labels" )
            __snake_case				=		model_inputs.pop("end_labels" )
            __snake_case				=		model_inputs.pop("pooled_labels" )

            __snake_case				=		state.apply_fn(**_lowerCAmelCase      ,	params=_lowerCAmelCase      ,	dropout_rng=_lowerCAmelCase      ,	train=_lowerCAmelCase )
            __snake_case		,					__snake_case		,					__snake_case				=		outputs

            return state.loss_fn(
                _lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	)

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

      __snake_case				=		state.apply_gradients(grads=_lowerCAmelCase )
      return state, metrics, new_drp_rng





@partial(jax.pmap      ,	axis_name="batch" )
def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	**_lowerCAmelCase )						->	List[str]:


      '''simple docstring'''






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

      __snake_case				=		state.apply_fn(**_lowerCAmelCase      ,	params=state.params      ,	train=_lowerCAmelCase )
      __snake_case		,					__snake_case		,					__snake_case				=		outputs

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






class     UpperCamelCase(  train_state.TrainState   ):
     snake_case_		:				Callable										=	struct.field(pytree_node=_a   )






@dataclass
class     UpperCamelCase:
     snake_case_		:				Args
     snake_case_		:				Callable
     snake_case_		:				Callable
     snake_case_		:				Callable
     snake_case_		:				Callable
     snake_case_		:				wandb
     snake_case_		:				Callable										=	None


     def 							SCREAMING_SNAKE_CASE_  (						self	:		str     ,    SCREAMING_SNAKE_CASE	:		Any     ,    SCREAMING_SNAKE_CASE	:		Tuple     ,    SCREAMING_SNAKE_CASE	:		Union[str, Any]     ,    SCREAMING_SNAKE_CASE	:		Tuple=None						)			->	Union[str, Any]:




           '''simple docstring'''

           __snake_case				=		model.params
           __snake_case				=		TrainState.create(
               apply_fn=model.__call__     ,    params=SCREAMING_SNAKE_CASE     ,    tx=SCREAMING_SNAKE_CASE     ,    loss_fn=SCREAMING_SNAKE_CASE     ,    )
           if ckpt_dir is not None:
                 __snake_case		,					__snake_case		,					__snake_case		,					__snake_case		,					__snake_case				=		restore_checkpoint(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE						)
                 __snake_case				=		{
                     "lr": args.lr,
                     "init_lr": args.init_lr,
                     "warmup_steps": args.warmup_steps,
                     "num_train_steps": num_train_steps,
                     "weight_decay": args.weight_decay,
                 }
                 __snake_case		,					__snake_case				=		build_tx(**SCREAMING_SNAKE_CASE						)
                 __snake_case				=		train_state.TrainState(
                     step=SCREAMING_SNAKE_CASE     ,    apply_fn=model.__call__     ,    params=SCREAMING_SNAKE_CASE     ,    tx=SCREAMING_SNAKE_CASE     ,    opt_state=SCREAMING_SNAKE_CASE     ,    )
                 __snake_case				=		args
                 __snake_case				=		data_collator
                 __snake_case				=		lr
                 __snake_case				=		params
           __snake_case				=		jax_utils.replicate(SCREAMING_SNAKE_CASE						)
           return state


     def 							SCREAMING_SNAKE_CASE_  (						self	:		Any     ,    SCREAMING_SNAKE_CASE	:		Optional[int]     ,    SCREAMING_SNAKE_CASE	:		Tuple     ,    SCREAMING_SNAKE_CASE	:		int						)			->	List[str]:




           '''simple docstring'''

           __snake_case				=		self.args
           __snake_case				=		len(SCREAMING_SNAKE_CASE						) // args.batch_size

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

                             __snake_case				=		self.evaluate(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE						)
                             __snake_case				=		{
                                 "step": state_step.item(),
                                 "eval_loss": eval_loss.item(),
                                 "tr_loss": tr_loss,
                                 "lr": lr.item(),
                             }
                             tqdm.write(str(SCREAMING_SNAKE_CASE						)						)
                             self.logger.log(SCREAMING_SNAKE_CASE     ,    commit=SCREAMING_SNAKE_CASE						)

                       if i % args.save_steps == 0:
                             self.save_checkpoint(args.save_dir + f'''-e{epoch}-s{i}'''     ,    state=SCREAMING_SNAKE_CASE						)


     def 							SCREAMING_SNAKE_CASE_  (						self	:		str     ,    SCREAMING_SNAKE_CASE	:		Any     ,    SCREAMING_SNAKE_CASE	:		List[str]						)			->	Any:




           '''simple docstring'''

           __snake_case				=		get_batched_dataset(SCREAMING_SNAKE_CASE     ,    self.args.batch_size						)
           __snake_case				=		len(SCREAMING_SNAKE_CASE						) // self.args.batch_size
           __snake_case				=		jnp.array(0     ,    dtype=jnp.floataa						)
           __snake_case				=		0
           for batch in tqdm(SCREAMING_SNAKE_CASE     ,    total=SCREAMING_SNAKE_CASE     ,    desc="Evaluating ... "						):
                 __snake_case				=		self.data_collator(SCREAMING_SNAKE_CASE						)
                 __snake_case				=		self.val_step_fn(SCREAMING_SNAKE_CASE     ,    **SCREAMING_SNAKE_CASE						)
                 running_loss += jax_utils.unreplicate(metrics["loss"]						)
                 i += 1
           return running_loss / i






     def 							SCREAMING_SNAKE_CASE_  (						self	:		int     ,    SCREAMING_SNAKE_CASE	:		int     ,    SCREAMING_SNAKE_CASE	:		Optional[int]						)			->	Any:




           '''simple docstring'''

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





def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	_lowerCAmelCase )						->	Any:


      '''simple docstring'''






      print(F'''RESTORING CHECKPOINT FROM {save_dir}'''      ,	end=" ... " )
      with open(os.path.join(_lowerCAmelCase      ,	"flax_model.msgpack" )      ,	"rb" ) as f:
            __snake_case				=		from_bytes(state.params      ,	f.read() )

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

      __snake_case				=		joblib.load(os.path.join(_lowerCAmelCase      ,	"args.joblib" ) )
      __snake_case				=		joblib.load(os.path.join(_lowerCAmelCase      ,	"data_collator.joblib" ) )

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

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





def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase )						->	str:


      '''simple docstring'''






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





def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase )						->	List[str]:


      '''simple docstring'''






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

      __snake_case				=		scheduler_fn(_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase )

      __snake_case				=		optax.adamw(learning_rate=_lowerCAmelCase      ,	weight_decay=_lowerCAmelCase      ,	mask=_lowerCAmelCase )
      return tx, lr









import os
import tempfile
import unittest
import uuid
from pathlib import Path

from transformers.testing_utils import get_tests_dir, require_soundfile, require_torch, require_vision
from transformers.tools.agent_types import AgentAudio, AgentImage, AgentText
from transformers.utils import is_soundfile_availble, is_torch_available, is_vision_available


if is_torch_available():
	import torch

if is_soundfile_availble():
	import soundfile as sf

if is_vision_available():
	from PIL import Image





def 	_lowerCAmelCase       (    _lowerCAmelCase="" )						->	str:


						'''simple docstring'''






						__snake_case				=		tempfile.mkdtemp()
						return os.path.join(_lowerCAmelCase      ,	str(uuid.uuida() ) + suffix )






@require_soundfile
@require_torch
class     UpperCamelCase(  unittest.TestCase   ):


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Union[str, Any]						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		torch.rand(1_2     ,    dtype=torch.floataa						) - 0.5
											__snake_case				=		AgentAudio(SCREAMING_SNAKE_CASE						)
											__snake_case				=		str(agent_type.to_string()						)

											# Ensure that the tensor and the agent_type's tensor are the same
											self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE     ,    agent_type.to_raw()     ,    atol=1e-4						)						)

											del agent_type

											# Ensure the path remains even after the object deletion
											self.assertTrue(os.path.exists(SCREAMING_SNAKE_CASE						)						)

											# Ensure that the file contains the same value as the original tensor
											__snake_case		,					__snake_case				=		sf.read(SCREAMING_SNAKE_CASE						)
											self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE     ,    torch.tensor(SCREAMING_SNAKE_CASE						)     ,    atol=1e-4						)						)






					def 							SCREAMING_SNAKE_CASE_  (						self	:		Dict						)			->	int:




											'''simple docstring'''

											__snake_case				=		torch.rand(1_2     ,    dtype=torch.floataa						) - 0.5
											__snake_case				=		get_new_path(suffix=".wav"						)
											sf.write(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    1_6_0_0_0						)

											__snake_case				=		AgentAudio(SCREAMING_SNAKE_CASE						)

											self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE     ,    agent_type.to_raw()     ,    atol=1e-4						)						)
											self.assertEqual(agent_type.to_string()     ,    SCREAMING_SNAKE_CASE						)






@require_vision
@require_torch
class     UpperCamelCase(  unittest.TestCase   ):


					def 							SCREAMING_SNAKE_CASE_  (						self	:		str						)			->	Optional[int]:




											'''simple docstring'''

											__snake_case				=		torch.randint(0     ,    2_5_6     ,    (6_4, 6_4, 3)						)
											__snake_case				=		AgentImage(SCREAMING_SNAKE_CASE						)
											__snake_case				=		str(agent_type.to_string()						)

											# Ensure that the tensor and the agent_type's tensor are the same
											self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE     ,    agent_type._tensor     ,    atol=1e-4						)						)

											self.assertIsInstance(agent_type.to_raw()     ,    Image.Image						)

											# Ensure the path remains even after the object deletion
											del agent_type
											self.assertTrue(os.path.exists(SCREAMING_SNAKE_CASE						)						)


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Tuple						)			->	Any:




											'''simple docstring'''

											__snake_case				=		Path(get_tests_dir("fixtures/tests_samples/COCO"						)						) / "000000039769.png"
											__snake_case				=		Image.open(SCREAMING_SNAKE_CASE						)
											__snake_case				=		AgentImage(SCREAMING_SNAKE_CASE						)

											self.assertTrue(path.samefile(agent_type.to_string()						)						)
											self.assertTrue(image == agent_type.to_raw()						)

											# Ensure the path remains even after the object deletion
											del agent_type
											self.assertTrue(os.path.exists(SCREAMING_SNAKE_CASE						)						)






					def 							SCREAMING_SNAKE_CASE_  (						self	:		Union[str, Any]						)			->	Union[str, Any]:




											'''simple docstring'''

											__snake_case				=		Path(get_tests_dir("fixtures/tests_samples/COCO"						)						) / "000000039769.png"
											__snake_case				=		Image.open(SCREAMING_SNAKE_CASE						)
											__snake_case				=		AgentImage(SCREAMING_SNAKE_CASE						)

											self.assertFalse(path.samefile(agent_type.to_string()						)						)
											self.assertTrue(image == agent_type.to_raw()						)

											# Ensure the path remains even after the object deletion
											del agent_type
											self.assertTrue(os.path.exists(SCREAMING_SNAKE_CASE						)						)






class     UpperCamelCase(  unittest.TestCase   ):


					def 							SCREAMING_SNAKE_CASE_  (						self	:		Any						)			->	Any:




											'''simple docstring'''

											__snake_case				=		"Hey!"
											__snake_case				=		AgentText(SCREAMING_SNAKE_CASE						)

											self.assertEqual(SCREAMING_SNAKE_CASE     ,    agent_type.to_string()						)
											self.assertEqual(SCREAMING_SNAKE_CASE     ,    agent_type.to_raw()						)
											self.assertEqual(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE						)





