

'''simple docstring'''
import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import MaskFormerForInstanceSegmentation, MaskFormerModel

     if is_vision_available():
          from transformers import MaskFormerImageProcessor

if is_vision_available():
     from PIL import Image

class 						__magic_name__   :







 def __init__(       self  , snake_case_  , snake_case_=2  , snake_case_=True  , snake_case_=False  , snake_case_=10  , snake_case_=3  , snake_case_=32 * 4  , snake_case_=32 * 6  , snake_case_=4  , snake_case_=32  , ):
   lowercase									=parent
   lowercase									=batch_size
   lowercase									=is_training
   lowercase									=use_auxiliary_loss
   lowercase									=num_queries
   lowercase									=num_channels
   lowercase									=min_size
   lowercase									=max_size
   lowercase									=num_labels
   lowercase									=mask_feature_size







 def 						_A(       self   ):
   lowercase									=floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]   ).to(
       snake_case_   )

   lowercase									=torch.ones([self.batch_size, self.min_size, self.max_size]  , device=snake_case_   )

   lowercase									=(
       torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]  , device=snake_case_   ) > 0.5
   ).float()
   lowercase									=(torch.rand((self.batch_size, self.num_labels)  , device=snake_case_   ) > 0.5).long()

   lowercase									=self.get_config()
   return config, pixel_values, pixel_mask, mask_labels, class_labels







 def 						_A(       self   ):
   return MaskFormerConfig.from_backbone_and_decoder_configs(
       backbone_config=SwinConfig(
           depths=[1, 1, 1, 1]  , )  , decoder_config=DetrConfig(
           decoder_ffn_dim=1_28  , num_queries=self.num_queries  , decoder_attention_heads=2  , d_model=self.mask_feature_size  , )  , mask_feature_size=self.mask_feature_size  , fpn_feature_size=self.mask_feature_size  , num_channels=self.num_channels  , num_labels=self.num_labels  , )







 def 						_A(       self   ):
   lowercase							, lowercase							, lowercase							, lowercase							, lowercase									=self.prepare_config_and_inputs()
   lowercase									={'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
   return config, inputs_dict







 def 						_A(       self  , snake_case_  , snake_case_   ):
   lowercase									=output.encoder_hidden_states
   lowercase									=output.pixel_decoder_hidden_states
   lowercase									=output.transformer_decoder_hidden_states

   self.parent.assertTrue(len(snake_case_   )  , len(config.backbone_config.depths   )   )
   self.parent.assertTrue(len(snake_case_   )  , len(config.backbone_config.depths   )   )
   self.parent.assertTrue(len(snake_case_   )  , config.decoder_config.decoder_layers   )







 def 						_A(       self  , snake_case_  , snake_case_  , snake_case_  , snake_case_=False   ):
   with torch.no_grad():
     lowercase									=MaskFormerModel(config=snake_case_   )
     model.to(snake_case_   )
     model.eval()

     lowercase									=model(pixel_values=snake_case_  , pixel_mask=snake_case_   )
     lowercase									=model(snake_case_  , output_hidden_states=snake_case_   )
   # the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
   # encoder and pixel decoder
   self.parent.assertEqual(
       output.transformer_decoder_last_hidden_state.shape  , (self.batch_size, self.num_queries, self.mask_feature_size)  , )
   # let's ensure the other two hidden state exists
   self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None   )
   self.parent.assertTrue(output.encoder_last_hidden_state is not None   )

   if output_hidden_states:
     self.check_output_hidden_state(snake_case_  , snake_case_   )







 def 						_A(       self  , snake_case_  , snake_case_  , snake_case_  , snake_case_  , snake_case_   ):
   lowercase									=MaskFormerForInstanceSegmentation(config=snake_case_   )
   model.to(snake_case_   )
   model.eval()

   def comm_check_on_output(snake_case_   ):
     # let's still check that all the required stuff is there
     self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None   )
     self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None   )
     self.parent.assertTrue(result.encoder_last_hidden_state is not None   )
     # okay, now we need to check the logits shape
     # due to the encoder compression, masks have a //4 spatial size
     self.parent.assertEqual(
         result.masks_queries_logits.shape  , (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)  , )
     # + 1 for null class
     self.parent.assertEqual(
         result.class_queries_logits.shape  , (self.batch_size, self.num_queries, self.num_labels + 1)   )

   with torch.no_grad():
     lowercase									=model(pixel_values=snake_case_  , pixel_mask=snake_case_   )
     lowercase									=model(snake_case_   )

     comm_check_on_output(snake_case_   )

     lowercase									=model(
         pixel_values=snake_case_  , pixel_mask=snake_case_  , mask_labels=snake_case_  , class_labels=snake_case_   )

   comm_check_on_output(snake_case_   )

   self.parent.assertTrue(result.loss is not None   )
   self.parent.assertEqual(result.loss.shape  , torch.Size([1]   )   )

@require_torch
class 						__magic_name__   (							__SCREAMING_SNAKE_CASE ,				__SCREAMING_SNAKE_CASE ,				unittest.TestCase ):
 UpperCamelCase__     =       (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
 UpperCamelCase__     =       (
     {'feature-extraction': MaskFormerModel, 'image-segmentation': MaskFormerForInstanceSegmentation}
     if is_torch_available()
     else {}
 )

 UpperCamelCase__     =       False
 UpperCamelCase__     =       False
 UpperCamelCase__     =       False
 UpperCamelCase__     =       False







 def 						_A(       self   ):
   lowercase									=MaskFormerModelTester(self   )
   lowercase									=ConfigTester(self  , config_class=snake_case_  , has_text_modality=snake_case_   )







 def 						_A(       self   ):
   self.config_tester.run_common_tests()







 def 						_A(       self   ):
   lowercase							, lowercase									=self.model_tester.prepare_config_and_inputs_for_common()
   self.model_tester.create_and_check_maskformer_model(snake_case_  , **snake_case_  , output_hidden_states=snake_case_   )







 def 						_A(       self   ):
   lowercase									=self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*snake_case_   )







 @unittest.skip(reason='''MaskFormer does not use inputs_embeds'''   )
 def 						_A(       self   ):
   pass







 @unittest.skip(reason='''MaskFormer does not have a get_input_embeddings method'''   )
 def 						_A(       self   ):
   pass







 @unittest.skip(reason='''MaskFormer is not a generative model'''   )
 def 						_A(       self   ):
   pass







 @unittest.skip(reason='''MaskFormer does not use token embeddings'''   )
 def 						_A(       self   ):
   pass







 @require_torch_multi_gpu
 @unittest.skip(
     reason='''MaskFormer has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''   )
 def 						_A(       self   ):
   pass







 @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''   )
 def 						_A(       self   ):
   pass







 def 						_A(       self   ):
   lowercase							, lowercase									=self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     lowercase									=model_class(snake_case_   )
     lowercase									=inspect.signature(model.forward   )
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     lowercase									=[*signature.parameters.keys()]

     lowercase									=['''pixel_values''']
     self.assertListEqual(arg_names[:1]  , snake_case_   )







 @slow
 def 						_A(       self   ):
   for model_name in ["facebook/maskformer-swin-small-coco"]:
     lowercase									=MaskFormerModel.from_pretrained(snake_case_   )
     self.assertIsNotNone(snake_case_   )







 def 						_A(       self   ):
   lowercase									=(self.model_tester.min_size,) * 2
   lowercase									={
       '''pixel_values''': torch.randn((2, 3, *size)  , device=snake_case_   ),
       '''mask_labels''': torch.randn((2, 10, *size)  , device=snake_case_   ),
       '''class_labels''': torch.zeros(2  , 10  , device=snake_case_   ).long(),
   }

   lowercase									=MaskFormerForInstanceSegmentation(MaskFormerConfig()   ).to(snake_case_   )
   lowercase									=model(**snake_case_   )
   self.assertTrue(outputs.loss is not None   )







 def 						_A(       self   ):
   lowercase							, lowercase									=self.model_tester.prepare_config_and_inputs_for_common()
   self.model_tester.create_and_check_maskformer_model(snake_case_  , **snake_case_  , output_hidden_states=snake_case_   )







 def 						_A(       self   ):
   lowercase							, lowercase									=self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     lowercase									=model_class(snake_case_   ).to(snake_case_   )
     lowercase									=model(**snake_case_  , output_attentions=snake_case_   )
     self.assertTrue(outputs.attentions is not None   )







 def 						_A(       self   ):
   if not self.model_tester.is_training:
     return
   # only MaskFormerForInstanceSegmentation has the loss
   lowercase									=self.all_model_classes[1]
   lowercase							, lowercase							, lowercase							, lowercase							, lowercase									=self.model_tester.prepare_config_and_inputs()

   lowercase									=model_class(snake_case_   )
   model.to(snake_case_   )
   model.train()

   lowercase									=model(snake_case_  , mask_labels=snake_case_  , class_labels=snake_case_   ).loss
   loss.backward()







 def 						_A(       self   ):
   # only MaskFormerForInstanceSegmentation has the loss
   lowercase									=self.all_model_classes[1]
   lowercase							, lowercase							, lowercase							, lowercase							, lowercase									=self.model_tester.prepare_config_and_inputs()
   lowercase									=True
   lowercase									=True

   lowercase									=model_class(snake_case_   )
   model.to(snake_case_   )
   model.train()

   lowercase									=model(snake_case_  , mask_labels=snake_case_  , class_labels=snake_case_   )

   lowercase									=outputs.encoder_hidden_states[0]
   encoder_hidden_states.retain_grad()

   lowercase									=outputs.pixel_decoder_hidden_states[0]
   pixel_decoder_hidden_states.retain_grad()
   # we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
   lowercase									=outputs.transformer_decoder_hidden_states[0]
   transformer_decoder_hidden_states.retain_grad()

   lowercase									=outputs.attentions[0]
   attentions.retain_grad()

   outputs.loss.backward(retain_graph=snake_case_   )

   self.assertIsNotNone(encoder_hidden_states.grad   )
   self.assertIsNotNone(pixel_decoder_hidden_states.grad   )
   self.assertIsNotNone(transformer_decoder_hidden_states.grad   )
   self.assertIsNotNone(attentions.grad   )

_UpperCAmelCase  :		List[Any]  =		1e-4




def 	UpperCamelCase   (       )		->	Optional[Any]:




  '''simple docstring'''



  lowercase									=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''	)
  return image




@require_vision
@slow
class 						__magic_name__   (							unittest.TestCase ):







 @cached_property
 def 						_A(       self   ):
   return (
       MaskFormerImageProcessor.from_pretrained('''facebook/maskformer-swin-small-coco'''   )
       if is_vision_available()
       else None
   )







 def 						_A(       self   ):
   lowercase									=MaskFormerModel.from_pretrained('''facebook/maskformer-swin-small-coco'''   ).to(snake_case_   )
   lowercase									=self.default_image_processor
   lowercase									=prepare_img()
   lowercase									=image_processor(snake_case_  , return_tensors='''pt'''   ).to(snake_case_   )
   lowercase									=inputs['''pixel_values'''].shape
   # check size is divisible by 32
   self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
   # check size
   self.assertEqual(snake_case_  , (1, 3, 8_00, 10_88)   )

   with torch.no_grad():
     lowercase									=model(**snake_case_   )

   lowercase									=torch.tensor(
       [[-0.04_82, 0.92_28, 0.49_51], [-0.25_47, 0.80_17, 0.85_27], [-0.00_69, 0.33_85, -0.00_89]]   ).to(snake_case_   )
   self.assertTrue(
       torch.allclose(
           outputs.encoder_last_hidden_state[0, 0, :3, :3]  , snake_case_  , atol=snake_case_   )   )

   lowercase									=torch.tensor(
       [[-0.84_22, -0.84_34, -0.97_18], [-1.01_44, -0.55_65, -0.41_95], [-1.00_38, -0.44_84, -0.19_61]]   ).to(snake_case_   )
   self.assertTrue(
       torch.allclose(
           outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]  , snake_case_  , atol=snake_case_   )   )

   lowercase									=torch.tensor(
       [[0.28_52, -0.01_59, 0.97_35], [0.62_54, 0.18_58, 0.85_29], [-0.06_80, -0.41_16, 1.84_13]]   ).to(snake_case_   )
   self.assertTrue(
       torch.allclose(
           outputs.transformer_decoder_last_hidden_state[0, :3, :3]  , snake_case_  , atol=snake_case_   )   )







 def 						_A(       self   ):
   lowercase									=(
       MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-swin-small-coco'''   )
       .to(snake_case_   )
       .eval()
   )
   lowercase									=self.default_image_processor
   lowercase									=prepare_img()
   lowercase									=image_processor(snake_case_  , return_tensors='''pt'''   ).to(snake_case_   )
   lowercase									=inputs['''pixel_values'''].shape
   # check size is divisible by 32
   self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
   # check size
   self.assertEqual(snake_case_  , (1, 3, 8_00, 10_88)   )

   with torch.no_grad():
     lowercase									=model(**snake_case_   )
   # masks_queries_logits
   lowercase									=outputs.masks_queries_logits
   self.assertEqual(
       masks_queries_logits.shape  , (1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)  , )
   lowercase									=[
       [-1.3_73_71_24, -1.7_72_49_37, -1.9_36_42_33],
       [-1.5_97_72_81, -1.9_86_79_39, -2.1_52_36_95],
       [-1.5_79_53_98, -1.9_26_98_32, -2.09_39_42],
   ]
   lowercase									=torch.tensor(snake_case_   ).to(snake_case_   )
   self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]  , snake_case_  , atol=snake_case_   )   )
   # class_queries_logits
   lowercase									=outputs.class_queries_logits
   self.assertEqual(
       class_queries_logits.shape  , (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)   )
   lowercase									=torch.tensor(
       [
           [1.6_512E00, -5.2_572E00, -3.3_519E00],
           [3.6_169E-02, -5.9_025E00, -2.9_313E00],
           [1.0_766E-04, -7.7_630E00, -5.1_263E00],
       ]   ).to(snake_case_   )
   self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]  , snake_case_  , atol=snake_case_   )   )







 def 						_A(       self   ):
   lowercase									=(
       MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-resnet101-coco-stuff'''   )
       .to(snake_case_   )
       .eval()
   )
   lowercase									=self.default_image_processor
   lowercase									=prepare_img()
   lowercase									=image_processor(snake_case_  , return_tensors='''pt'''   ).to(snake_case_   )
   lowercase									=inputs['''pixel_values'''].shape
   # check size is divisible by 32
   self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
   # check size
   self.assertEqual(snake_case_  , (1, 3, 8_00, 10_88)   )

   with torch.no_grad():
     lowercase									=model(**snake_case_   )
   # masks_queries_logits
   lowercase									=outputs.masks_queries_logits
   self.assertEqual(
       masks_queries_logits.shape  , (1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)  , )
   lowercase									=[[-0.90_46, -2.63_66, -4.60_62], [-3.41_79, -5.78_90, -8.80_57], [-4.91_79, -7.65_60, -10.77_11]]
   lowercase									=torch.tensor(snake_case_   ).to(snake_case_   )
   self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]  , snake_case_  , atol=snake_case_   )   )
   # class_queries_logits
   lowercase									=outputs.class_queries_logits
   self.assertEqual(
       class_queries_logits.shape  , (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)   )
   lowercase									=torch.tensor(
       [[4.71_88, -3.25_85, -2.88_57], [6.68_71, -2.91_81, -1.24_87], [7.24_49, -2.27_64, -2.18_74]]   ).to(snake_case_   )
   self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]  , snake_case_  , atol=snake_case_   )   )







 def 						_A(       self   ):
   lowercase									=(
       MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-swin-small-coco'''   )
       .to(snake_case_   )
       .eval()
   )
   lowercase									=self.default_image_processor

   lowercase									=image_processor(
       [np.zeros((3, 8_00, 13_33)   ), np.zeros((3, 8_00, 13_33)   )]  , segmentation_maps=[np.zeros((3_84, 3_84)   ).astype(np.floataa   ), np.zeros((3_84, 3_84)   ).astype(np.floataa   )]  , return_tensors='''pt'''  , )

   lowercase									=inputs['''pixel_values'''].to(snake_case_   )
   lowercase									=[el.to(snake_case_   ) for el in inputs['''mask_labels''']]
   lowercase									=[el.to(snake_case_   ) for el in inputs['''class_labels''']]

   with torch.no_grad():
     lowercase									=model(**snake_case_   )

   self.assertTrue(outputs.loss is not None   )







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_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       42
							lowerCamelCase_							     =       jnp.floataa
							lowerCamelCase_							     =       True


							def   _snake_case	(				self     :Tuple		) -> Optional[Any]:



									"""simple docstring"""



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


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



									"""simple docstring"""



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




class 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       FlaxBigBirdForNaturalQuestionsModule






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

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




@dataclass
class 	UpperCamelCase_							:
							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   _snake_case	(				self     :str		) -> Optional[int]:



									"""simple docstring"""



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




@dataclass
class 	UpperCamelCase_							:
							lowerCamelCase_							     =       42
							lowerCamelCase_							     =       40_96  # no dynamic padding on TPUs


							def __call__(				self     :Optional[Any]				,					__A     :Optional[int]		) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.collate_fn(__A		)
									SCREAMING_SNAKE_CASE__								=   jax.tree_util.tree_map(__A				,					__A		)
									return batch


							def   _snake_case	(				self     :List[Any]				,					__A     :Union[str, Any]		) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   self.fetch_inputs(features["""input_ids"""]		)
									SCREAMING_SNAKE_CASE__								=   {
									    """input_ids""": jnp.array(__A				,					dtype=jnp.intaa		),
									    """attention_mask""": jnp.array(__A				,					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   _snake_case	(				self     :Tuple				,					__A     :list		) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   [self._fetch_inputs(__A		) for ids in input_ids]
									return zip(*__A		)


							def   _snake_case	(				self     :List[str]				,					__A     :list		) -> Union[str, Any]:



									"""simple docstring"""



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






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				str       , UpperCamelCase__:				List[str]       , UpperCamelCase__:				Optional[Any]=None     ):
		if seed is not None:
				SCREAMING_SNAKE_CASE__								=   dataset.shuffle(seed=UpperCamelCase__     )
		for i in range(len(UpperCamelCase__     ) // batch_size     ):
				SCREAMING_SNAKE_CASE__								=   dataset[i * batch_size : (i + 1) * batch_size]
				yield dict(UpperCamelCase__     )






@partial(jax.pmap       , axis_name="""batch"""     )
def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				Dict       , UpperCamelCase__:				Optional[int]       , **UpperCamelCase__:				Optional[int]     ):
		def loss_fn(UpperCamelCase__:				List[Any]     ):
				SCREAMING_SNAKE_CASE__								=   model_inputs.pop("""start_labels"""     )
				SCREAMING_SNAKE_CASE__								=   model_inputs.pop("""end_labels"""     )
				SCREAMING_SNAKE_CASE__								=   model_inputs.pop("""pooled_labels"""     )

				SCREAMING_SNAKE_CASE__								=   state.apply_fn(**UpperCamelCase__       , params=UpperCamelCase__       , dropout_rng=UpperCamelCase__       , train=UpperCamelCase__     )
				SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   outputs

				return state.loss_fn(
				    UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , )

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

		SCREAMING_SNAKE_CASE__								=   state.apply_gradients(grads=UpperCamelCase__     )
		return state, metrics, new_drp_rng






@partial(jax.pmap       , axis_name="""batch"""     )
def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				Optional[Any]       , **UpperCamelCase__:				Dict     ):
		SCREAMING_SNAKE_CASE__								=   model_inputs.pop("""start_labels"""     )
		SCREAMING_SNAKE_CASE__								=   model_inputs.pop("""end_labels"""     )
		SCREAMING_SNAKE_CASE__								=   model_inputs.pop("""pooled_labels"""     )

		SCREAMING_SNAKE_CASE__								=   state.apply_fn(**UpperCamelCase__       , params=state.params       , train=UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   outputs

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




class 	UpperCamelCase_							(  train_state.TrainState     ):
							lowerCamelCase_							     =       struct.field(pytree_node=UpperCamelCase__     )




@dataclass
class 	UpperCamelCase_							:
							lowerCamelCase_							     =       42
							lowerCamelCase_							     =       42
							lowerCamelCase_							     =       42
							lowerCamelCase_							     =       42
							lowerCamelCase_							     =       42
							lowerCamelCase_							     =       42
							lowerCamelCase_							     =       None


							def   _snake_case	(				self     :List[Any]				,					__A     :str				,					__A     :str				,					__A     :str				,					__A     :Tuple=None		) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   model.params
									SCREAMING_SNAKE_CASE__								=   TrainState.create(
									    apply_fn=model.__call__				,					params=__A				,					tx=__A				,					loss_fn=__A				,					)
									if ckpt_dir is not None:
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   restore_checkpoint(__A				,					__A		)
											SCREAMING_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,
											}
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   build_tx(**__A		)
											SCREAMING_SNAKE_CASE__								=   train_state.TrainState(
											    step=__A				,					apply_fn=model.__call__				,					params=__A				,					tx=__A				,					opt_state=__A				,					)
											SCREAMING_SNAKE_CASE__								=   args
											SCREAMING_SNAKE_CASE__								=   data_collator
											SCREAMING_SNAKE_CASE__								=   lr
											SCREAMING_SNAKE_CASE__								=   params
									SCREAMING_SNAKE_CASE__								=   jax_utils.replicate(__A		)
									return state


							def   _snake_case	(				self     :Optional[Any]				,					__A     :Optional[int]				,					__A     :int				,					__A     :int		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.args
									SCREAMING_SNAKE_CASE__								=   len(__A		) // args.batch_size

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

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

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


							def   _snake_case	(				self     :List[str]				,					__A     :Dict				,					__A     :str		) -> Union[str, Any]:



									"""simple docstring"""



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


							def   _snake_case	(				self     :List[Any]				,					__A     :Any				,					__A     :Dict		) -> Union[str, Any]:



									"""simple docstring"""



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






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				int       , UpperCamelCase__:				Optional[Any]     ):
		print(f'''RESTORING CHECKPOINT FROM {save_dir}'''       , end=""" ... """     )
		with open(os.path.join(UpperCamelCase__       , """flax_model.msgpack"""     )       , """rb"""     ) as f:
				SCREAMING_SNAKE_CASE__								=   from_bytes(state.params       , f.read()     )

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

		SCREAMING_SNAKE_CASE__								=   joblib.load(os.path.join(UpperCamelCase__       , """args.joblib"""     )     )
		SCREAMING_SNAKE_CASE__								=   joblib.load(os.path.join(UpperCamelCase__       , """data_collator.joblib"""     )     )

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

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






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				int       , UpperCamelCase__:				Union[str, Any]       , UpperCamelCase__:				List[Any]       , UpperCamelCase__:				Dict     ):
		SCREAMING_SNAKE_CASE__								=   num_train_steps - warmup_steps
		SCREAMING_SNAKE_CASE__								=   optax.linear_schedule(init_value=UpperCamelCase__       , end_value=UpperCamelCase__       , transition_steps=UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__								=   optax.linear_schedule(init_value=UpperCamelCase__       , end_value=1e-7       , transition_steps=UpperCamelCase__     )
		SCREAMING_SNAKE_CASE__								=   optax.join_schedules(schedules=[warmup_fn, decay_fn]       , boundaries=[warmup_steps]     )
		return lr






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				List[str]       , UpperCamelCase__:				Union[str, Any]       , UpperCamelCase__:				Tuple       , UpperCamelCase__:				Tuple       , UpperCamelCase__:				Tuple     ):
		def weight_decay_mask(UpperCamelCase__:				Any     ):
				SCREAMING_SNAKE_CASE__								=   traverse_util.flatten_dict(UpperCamelCase__     )
				SCREAMING_SNAKE_CASE__								=   {k: (v[-1] != """bias""" and v[-2:] != ("""LayerNorm""", """scale""")) for k, v in params.items()}
				return traverse_util.unflatten_dict(UpperCamelCase__     )

		SCREAMING_SNAKE_CASE__								=   scheduler_fn(UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__       , UpperCamelCase__     )

		SCREAMING_SNAKE_CASE__								=   optax.adamw(learning_rate=UpperCamelCase__       , weight_decay=UpperCamelCase__       , mask=UpperCamelCase__     )
		return tx, lr