






from __future__ import annotations

import unittest

from transformers import CTRLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers.models.ctrl.modeling_tf_ctrl import (
       TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFCTRLForSequenceClassification,
       TFCTRLLMHeadModel,
       TFCTRLModel,
   )


class 		lowerCAmelCase_							(     UpperCamelCase   ):





     def __init__(  self       ,   _UpperCAmelCase       ,   )		->							List[str]:
        lowerCamelCase											=parent
        lowerCamelCase											=13
        lowerCamelCase											=7
        lowerCamelCase											=True
        lowerCamelCase											=True
        lowerCamelCase											=True
        lowerCamelCase											=True
        lowerCamelCase											=True
        lowerCamelCase											=99
        lowerCamelCase											=32
        lowerCamelCase											=2
        lowerCamelCase											=4
        lowerCamelCase											=37
        lowerCamelCase											="""gelu"""
        lowerCamelCase											=0.1
        lowerCamelCase											=0.1
        lowerCamelCase											=512
        lowerCamelCase											=16
        lowerCamelCase											=2
        lowerCamelCase											=0.02
        lowerCamelCase											=3
        lowerCamelCase											=4
        lowerCamelCase											=None
        lowerCamelCase											=self.vocab_size - 1





     def 	lowerCamelCase       (  self     )		->							int:
        lowerCamelCase											=ids_tensor([self.batch_size, self.seq_length]       ,   self.vocab_size     )

        lowerCamelCase											=None
        if self.use_input_mask:
           lowerCamelCase											=random_attention_mask([self.batch_size, self.seq_length]     )

        lowerCamelCase											=None
        if self.use_token_type_ids:
           lowerCamelCase											=ids_tensor([self.batch_size, self.seq_length]       ,   self.type_vocab_size     )

        lowerCamelCase											=None
        if self.use_mc_token_ids:
           lowerCamelCase											=ids_tensor([self.batch_size, self.num_choices]       ,   self.seq_length     )

        lowerCamelCase											=None
        lowerCamelCase											=None
        lowerCamelCase											=None
        if self.use_labels:
           lowerCamelCase											=ids_tensor([self.batch_size]       ,   self.type_sequence_label_size     )
           lowerCamelCase											=ids_tensor([self.batch_size, self.seq_length]       ,   self.num_labels     )
           lowerCamelCase											=ids_tensor([self.batch_size]       ,   self.num_choices     )

        lowerCamelCase											=CTRLConfig(
            vocab_size=self.vocab_size       ,   n_embd=self.hidden_size       ,   n_layer=self.num_hidden_layers       ,   n_head=self.num_attention_heads       ,   n_positions=self.max_position_embeddings       ,   pad_token_id=self.pad_token_id       ,   )

        lowerCamelCase											=ids_tensor([self.num_hidden_layers, self.num_attention_heads]       ,   2     )

        return (
            config,
            input_ids,
            input_mask,
            head_mask,
            token_type_ids,
            mc_token_ids,
            sequence_labels,
            token_labels,
            choice_labels,
        )





     def 	lowerCamelCase       (  self       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   *_UpperCAmelCase     )		->							Union[str, Any]:
        lowerCamelCase											=TFCTRLModel(config=_UpperCAmelCase     )
        lowerCamelCase											={"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
        lowerCamelCase											=model(_UpperCAmelCase     )

        lowerCamelCase											=[input_ids, None, input_mask]  # None is the input for 'past'
        lowerCamelCase											=model(_UpperCAmelCase     )

        lowerCamelCase											=model(_UpperCAmelCase     )

        self.parent.assertEqual(result.last_hidden_state.shape       ,   (self.batch_size, self.seq_length, self.hidden_size)     )





     def 	lowerCamelCase       (  self       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   *_UpperCAmelCase     )		->							Optional[int]:
        lowerCamelCase											=TFCTRLLMHeadModel(config=_UpperCAmelCase     )
        lowerCamelCase											={"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
        lowerCamelCase											=model(_UpperCAmelCase     )
        self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.vocab_size)     )





     def 	lowerCamelCase       (  self       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   *_UpperCAmelCase     )		->							List[str]:
        lowerCamelCase											=self.num_labels
        lowerCamelCase											=ids_tensor([self.batch_size]       ,   self.type_sequence_label_size     )
        lowerCamelCase											={
            """input_ids""": input_ids,
            """token_type_ids""": token_type_ids,
            """labels""": sequence_labels,
        }
        lowerCamelCase											=TFCTRLForSequenceClassification(_UpperCAmelCase     )
        lowerCamelCase											=model(_UpperCAmelCase     )
        self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.num_labels)     )





     def 	lowerCamelCase       (  self     )		->							str:
        lowerCamelCase											=self.prepare_config_and_inputs()

        (
            (
            lowerCamelCase
        )							,					(
            lowerCamelCase
        )							,					(
            lowerCamelCase
        )							,					(
            lowerCamelCase
        )							,					(
            lowerCamelCase
        )							,					(
            lowerCamelCase
        )							,					(
            lowerCamelCase
        )							,					(
            lowerCamelCase
        )							,					(
            lowerCamelCase
        )							,					
        )											=config_and_inputs

        lowerCamelCase											={"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
        return config, inputs_dict




@require_tf
class 		lowerCAmelCase_							(     UpperCamelCase							,    UpperCamelCase							,    unittest.TestCase   ):
     __lowerCamelCase					:						Union[str, Any]     =    (TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
     __lowerCamelCase					:						int     =    (TFCTRLLMHeadModel,) if is_tf_available() else ()
     __lowerCamelCase					:						Any     =    (
         {
             '''feature-extraction''': TFCTRLModel,
             '''text-classification''': TFCTRLForSequenceClassification,
             '''text-generation''': TFCTRLLMHeadModel,
             '''zero-shot''': TFCTRLForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     __lowerCamelCase					:						int     =    False
     __lowerCamelCase					:						List[Any]     =    False





     def 	lowerCamelCase       (  self       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase     )		->							Any:
        if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
           # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
           # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
           # config could not be created.
           return True

        return False





     def 	lowerCamelCase       (  self     )		->							Tuple:
        lowerCamelCase											=TFCTRLModelTester(self     )
        lowerCamelCase											=ConfigTester(self       ,   config_class=_UpperCAmelCase       ,   n_embd=37     )





     def 	lowerCamelCase       (  self     )		->							List[Any]:
        self.config_tester.run_common_tests()





     def 	lowerCamelCase       (  self     )		->							Optional[int]:
        lowerCamelCase											=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_ctrl_model(*_UpperCAmelCase     )





     def 	lowerCamelCase       (  self     )		->							List[str]:
        lowerCamelCase											=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_ctrl_lm_head(*_UpperCAmelCase     )





     def 	lowerCamelCase       (  self     )		->							Tuple:
        lowerCamelCase											=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_ctrl_for_sequence_classification(*_UpperCAmelCase     )





     def 	lowerCamelCase       (  self     )		->							List[str]:
        lowerCamelCase							,					lowerCamelCase											=self.model_tester.prepare_config_and_inputs_for_common()
        lowerCamelCase											=[TFCTRLLMHeadModel]
        lowerCamelCase											=[TFCTRLForSequenceClassification]

        for model_class in self.all_model_classes:
           lowerCamelCase											=model_class(_UpperCAmelCase     )
           model.build()  # may be needed for the get_bias() call below
           assert isinstance(model.get_input_embeddings()       ,   tf.keras.layers.Layer     )

           if model_class in list_lm_models:
              lowerCamelCase											=model.get_output_embeddings()
              assert isinstance(_UpperCAmelCase       ,   tf.keras.layers.Layer     )
              lowerCamelCase											=model.get_bias()
              assert isinstance(_UpperCAmelCase       ,   _UpperCAmelCase     )
              for k, v in name.items():
                 assert isinstance(_UpperCAmelCase       ,   tf.Variable     )
           elif model_class in list_other_models_with_output_ebd:
              lowerCamelCase											=model.get_output_embeddings()
              assert isinstance(_UpperCAmelCase       ,   tf.keras.layers.Layer     )
              lowerCamelCase											=model.get_bias()
              assert name is None
           else:
              lowerCamelCase											=model.get_output_embeddings()
              assert x is None
              lowerCamelCase											=model.get_bias()
              assert name is None





     @slow
     def 	lowerCamelCase       (  self     )		->							Union[str, Any]:
        for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowerCamelCase											=TFCTRLModel.from_pretrained(_UpperCAmelCase     )
           self.assertIsNotNone(_UpperCAmelCase     )




@require_tf
class 		lowerCAmelCase_							(     unittest.TestCase   ):





     @slow
     def 	lowerCamelCase       (  self     )		->							int:
        lowerCamelCase											=TFCTRLLMHeadModel.from_pretrained("""ctrl"""     )
        lowerCamelCase											=tf.convert_to_tensor([[11859, 0, 1611, 8]]       ,   dtype=tf.intaa     )  # Legal the president is
        lowerCamelCase											=[
            11859,
            0,
            1611,
            8,
            5,
            150,
            26449,
            2,
            19,
            348,
            469,
            3,
            2595,
            48,
            20740,
            246533,
            246533,
            19,
            30,
            5,
        ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

        lowerCamelCase											=model.generate(_UpperCAmelCase       ,   do_sample=_UpperCAmelCase     )
        self.assertListEqual(output_ids[0].numpy().tolist()       ,   _UpperCAmelCase     )







import os
from dataclasses import replace

import jax
import wandb
from bigbird_flax import Args, DataCollator, FlaxBigBirdForNaturalQuestions, Trainer, build_tx, train_step, val_step
from datasets import load_dataset
from flax import jax_utils

from transformers import BigBirdTokenizerFast


if __name__ == "__main__":
   print('#################### AVAILABLE DEVICES ####################')
   print(jax.devices())
   print('###########################################################')

   # setup for wandb sweep
   _UpperCAmelCase				:	Optional[Any]    						=    Args()
   _UpperCAmelCase				:	Tuple    						=    wandb.init(project='bigbird-natural-questions', config=args.__dict__)
   _UpperCAmelCase				:	Optional[Any]    						=    dict(logger.config)
   del wandb_args["batch_size"]
   _UpperCAmelCase				:	str    						=    replace(args, **wandb_args)
   _UpperCAmelCase				:	Dict    						=    args.base_dir + '-' + wandb.run.id
   _UpperCAmelCase				:	Tuple    						=    replace(args, base_dir=base_dir)
   print(args)

   _UpperCAmelCase				:	List[Any]    						=    load_dataset('json', data_files=args.tr_data_path)['train']
   _UpperCAmelCase				:	Dict    						=    load_dataset('json', data_files=args.val_data_path)['train']

   # drop extra batch for now
   _UpperCAmelCase				:	List[str]    						=    range(len(tr_dataset) - len(tr_dataset) % args.batch_size)
   _UpperCAmelCase				:	List[Any]    						=    tr_dataset.shuffle().select(indices)
   _UpperCAmelCase				:	Any    						=    range(len(val_dataset) - len(val_dataset) % args.batch_size)
   _UpperCAmelCase				:	Union[str, Any]    						=    val_dataset.shuffle().select(indices)

   if os.environ.get('TRAIN_ON_SMALL', 'false') == "true":
      _UpperCAmelCase				:	Optional[Any]    						=    tr_dataset.shuffle().select(range(8_0000))
      _UpperCAmelCase				:	List[Any]    						=    val_dataset.shuffle().select(range(8000))

   print(tr_dataset)
   print(val_dataset)

   _UpperCAmelCase				:	int    						=    FlaxBigBirdForNaturalQuestions.from_pretrained(
       args.model_id, block_size=args.block_size, num_random_blocks=args.num_random_blocks
   )
   _UpperCAmelCase				:	Union[str, Any]    						=    BigBirdTokenizerFast.from_pretrained(args.model_id)
   _UpperCAmelCase				:	List[str]    						=    DataCollator(pad_id=tokenizer.pad_token_id, max_length=4096)

   _UpperCAmelCase				:	Any    						=    {
       'lr': args.lr,
       'init_lr': args.init_lr,
       'warmup_steps': args.warmup_steps,
       'num_train_steps': args.max_epochs * (len(tr_dataset) // args.batch_size),
       'weight_decay': args.weight_decay,
   }
   _UpperCAmelCase						,						_UpperCAmelCase				:	Optional[int]    						=    build_tx(**tx_args)

   _UpperCAmelCase				:	List[Any]    						=    Trainer(
       args=args,
       data_collator=data_collator,
       model_save_fn=model.save_pretrained,
       train_step_fn=train_step,
       val_step_fn=val_step,
       logger=logger,
       scheduler_fn=lr,
   )

   _UpperCAmelCase				:	Union[str, Any]    						=    None
   _UpperCAmelCase				:	Optional[Any]    						=    trainer.create_state(model, tx, num_train_steps=tx_args['num_train_steps'], ckpt_dir=ckpt_dir)
   try:
      trainer.train(state, tr_dataset, val_dataset)
   except KeyboardInterrupt:
      print('Oooops; TRAINING STOPPED UNFORTUNATELY')

   print('SAVING WEIGHTS IN `final-weights`')
   _UpperCAmelCase				:	Optional[int]    						=    jax_utils.unreplicate(state.params)
   model.save_pretrained(os.path.join(args.base_dir, 'final-weights'), params=params)