



import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM



@require_torch
class     _UpperCAmelCase			:




     def __init__(				self ,						a__ ,						a__=9_9 ,						a__=1_3 ,						a__=1_6 ,						a__=7 ,						a__=True ,						a__=True ,						a__=True ,						a__=False ,						a__=True ,						a__=2 ,						a__=3_2 ,						a__=4 ,						a__=4 ,						a__=3_0 ,						a__=0 ,						a__=1 ,						a__=2 ,						a__=None ,						):
         A__		 =					parent
         A__		 =					batch_size
         A__		 =					decoder_seq_length
         # For common tests
         A__		 =					self.decoder_seq_length
         A__		 =					is_training
         A__		 =					use_attention_mask
         A__		 =					use_labels

         A__		 =					vocab_size
         A__		 =					d_model
         A__		 =					d_model
         A__		 =					decoder_layers
         A__		 =					decoder_layers
         A__		 =					decoder_ffn_dim
         A__		 =					decoder_attention_heads
         A__		 =					decoder_attention_heads
         A__		 =					eos_token_id
         A__		 =					bos_token_id
         A__		 =					pad_token_id
         A__		 =					decoder_start_token_id
         A__		 =					use_cache
         A__		 =					max_position_embeddings

         A__		 =					None
         A__		 =					decoder_seq_length
         A__		 =					2
         A__		 =					1




     def 		snake_case_		(				self):
         A__		 =					ids_tensor([self.batch_size, self.decoder_seq_length] ,						self.vocab_size)

         A__		 =					None
         if self.use_attention_mask:
             A__		 =					ids_tensor([self.batch_size, self.decoder_seq_length] ,						vocab_size=2)

         A__		 =					None
         if self.use_labels:
             A__		 =					ids_tensor([self.batch_size, self.decoder_seq_length] ,						self.vocab_size)

         A__		 =					TrOCRConfig(
             vocab_size=self.vocab_size ,						d_model=self.d_model ,						decoder_layers=self.decoder_layers ,						decoder_ffn_dim=self.decoder_ffn_dim ,						decoder_attention_heads=self.decoder_attention_heads ,						eos_token_id=self.eos_token_id ,						bos_token_id=self.bos_token_id ,						use_cache=self.use_cache ,						pad_token_id=self.pad_token_id ,						decoder_start_token_id=self.decoder_start_token_id ,						max_position_embeddings=self.max_position_embeddings ,						)

         return (config, input_ids, attention_mask, lm_labels)




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						):
         A__		 =					True
         A__		 =					TrOCRDecoder(config=a__).to(a__).eval()
         A__		 =					input_ids[:2]

         input_ids[input_ids == 0] += 1
         # first forward pass
         A__		 =					model(a__ ,						use_cache=a__)
         A__		 =					model(a__)
         A__		 =					model(a__ ,						use_cache=a__)

         self.parent.assertTrue(len(a__) == len(a__))
         self.parent.assertTrue(len(a__) == len(a__) + 1)

         A__		 =					outputs['''past_key_values''']

         # create hypothetical next token and extent to next_input_ids
         A__		 =					ids_tensor((2, 1) ,						config.vocab_size - 1) + 1

         # append to next input_ids and
         A__		 =					torch.cat([input_ids, next_tokens] ,						dim=-1)

         A__		 =					model(a__)['''last_hidden_state''']
         A__		 =					model(a__ ,						past_key_values=a__)['''last_hidden_state''']

         # select random slice
         A__		 =					ids_tensor((1,) ,						output_from_past.shape[-1]).item()
         A__		 =					output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
         A__		 =					output_from_past[:, 0, random_slice_idx].detach()

         # test that outputs are equal for slice
         assert torch.allclose(a__ ,						a__ ,						atol=1e-3)




     def 		snake_case_		(				self):
         A__		 =					self.prepare_config_and_inputs()
         A__ ,	A__ ,	A__ ,	A__		 =					config_and_inputs

         A__		 =					{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
         return config, inputs_dict







@require_torch
class     _UpperCAmelCase			(       A__			,      A__			,      A__			,      unittest.TestCase     ):
     UpperCamelCase__		    =				(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
     UpperCamelCase__		    =				(TrOCRForCausalLM,) if is_torch_available() else ()
     UpperCamelCase__		    =				{'''text-generation''': TrOCRForCausalLM} if is_torch_available() else {}
     UpperCamelCase__		    =				True
     UpperCamelCase__		    =				False




     def 		snake_case_		(				self):
         A__		 =					TrOCRStandaloneDecoderModelTester(self ,						is_training=a__)
         A__		 =					ConfigTester(self ,						config_class=a__)




     def 		snake_case_		(				self):
         pass




     def 		snake_case_		(				self):
         pass




     def 		snake_case_		(				self):
         pass




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




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_past(*a__)




     def 		snake_case_		(				self):
         return




     @unittest.skip('''The model doesn\'t support left padding''')  # and it's not used enough to be worth fixing :)
     def 		snake_case_		(				self):
         pass





import argparse
import os

# New Code #
import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


_lowercase						  =		16
_lowercase						  =		32



def 		lowerCAmelCase__    (					UpperCamelCase_				:       Accelerator ,       UpperCamelCase_				:       int = 1_6	)->  List[str]:
    A__		 =					AutoTokenizer.from_pretrained('''bert-base-cased'''	)
    A__		 =					load_dataset('''glue''' ,       '''mrpc'''	)

    def tokenize_function(UpperCamelCase_				:       Any	):
        # max_length=None => use the model max length (it's actually the default)
        A__		 =					tokenizer(examples['''sentence1'''] ,       examples['''sentence2'''] ,       truncation=UpperCamelCase_ ,       max_length=UpperCamelCase_	)
        return outputs

    # Apply the method we just defined to all the examples in all the splits of the dataset
    # starting with the main process first:
    with accelerator.main_process_first():
        A__		 =					datasets.map(
            UpperCamelCase_ ,       batched=UpperCamelCase_ ,       remove_columns=['''idx''', '''sentence1''', '''sentence2'''] ,       )

    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
    # transformers library
    A__		 =					tokenized_datasets.rename_column('''label''' ,       '''labels'''	)

    def collate_fn(UpperCamelCase_				:       str	):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        A__		 =					1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
        # When using mixed precision we want round multiples of 8/16
        if accelerator.mixed_precision == "fp8":
            A__		 =					1_6
        elif accelerator.mixed_precision != "no":
            A__		 =					8
        else:
            A__		 =					None

        return tokenizer.pad(
            UpperCamelCase_ ,       padding='''longest''' ,       max_length=UpperCamelCase_ ,       pad_to_multiple_of=UpperCamelCase_ ,       return_tensors='''pt''' ,       )

    # Instantiate dataloaders.
    A__		 =					DataLoader(
        tokenized_datasets['''train'''] ,       shuffle=UpperCamelCase_ ,       collate_fn=UpperCamelCase_ ,       batch_size=UpperCamelCase_	)
    A__		 =					DataLoader(
        tokenized_datasets['''validation'''] ,       shuffle=UpperCamelCase_ ,       collate_fn=UpperCamelCase_ ,       batch_size=UpperCamelCase_	)

    return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
   from accelerate.test_utils.training import mocked_dataloaders

   _lowercase						  =		mocked_dataloaders  # noqa: F811



def 		lowerCAmelCase__    (					UpperCamelCase_				:       Optional[Any] ,       UpperCamelCase_				:       Tuple	)->  Union[str, Any]:
    # For testing only
    if os.environ.get('''TESTING_MOCKED_DATALOADERS''' ,       UpperCamelCase_	) == "1":
        A__		 =					2
    # Initialize accelerator
    A__		 =					Accelerator(cpu=args.cpu ,       mixed_precision=args.mixed_precision	)
    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    A__		 =					config['''lr''']
    A__		 =					int(config['''num_epochs''']	)
    A__		 =					int(config['''seed''']	)
    A__		 =					int(config['''batch_size''']	)

    A__		 =					evaluate.load('''glue''' ,       '''mrpc'''	)

    # New Code #
    # We now can define an inner training loop function. It should take a batch size as the only parameter,
    # and build the dataloaders in there.
    # It also gets our decorator
    @find_executable_batch_size(starting_batch_size=UpperCamelCase_	)
    def inner_training_loop(UpperCamelCase_				:       Optional[int]	):
        # And now just move everything below under this function
        # We need to bring in the Accelerator object from earlier
        nonlocal accelerator
        # And reset all of its attributes that could hold onto any memory:
        accelerator.free_memory()

        # Then we can declare the model, optimizer, and everything else:
        set_seed(UpperCamelCase_	)

        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
        A__		 =					AutoModelForSequenceClassification.from_pretrained('''bert-base-cased''' ,       return_dict=UpperCamelCase_	)

        # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
        # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
        # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
        A__		 =					model.to(accelerator.device	)

        # Instantiate optimizer
        A__		 =					AdamW(params=model.parameters() ,       lr=UpperCamelCase_	)
        A__ ,	A__		 =					get_dataloaders(UpperCamelCase_ ,       UpperCamelCase_	)

        # Instantiate scheduler
        A__		 =					get_linear_schedule_with_warmup(
            optimizer=UpperCamelCase_ ,       num_warmup_steps=1_0_0 ,       num_training_steps=(len(UpperCamelCase_	) * num_epochs) ,       )

        # Prepare everything
        # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
        # prepare method.
        A__ ,	A__ ,	A__ ,	A__ ,	A__		 =					accelerator.prepare(
            UpperCamelCase_ ,       UpperCamelCase_ ,       UpperCamelCase_ ,       UpperCamelCase_ ,       UpperCamelCase_	)

        # Now we train the model
        for epoch in range(UpperCamelCase_	):
            model.train()
            for step, batch in enumerate(UpperCamelCase_	):
                # We could avoid this line since we set the accelerator with `device_placement=True`.
                batch.to(accelerator.device	)
                A__		 =					model(**UpperCamelCase_	)
                A__		 =					outputs.loss
                accelerator.backward(UpperCamelCase_	)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            model.eval()
            for step, batch in enumerate(UpperCamelCase_	):
                # We could avoid this line since we set the accelerator with `device_placement=True`.
                batch.to(accelerator.device	)
                with torch.no_grad():
                    A__		 =					model(**UpperCamelCase_	)
                A__		 =					outputs.logits.argmax(dim=-1	)
                A__ ,	A__		 =					accelerator.gather_for_metrics((predictions, batch['''labels'''])	)
                metric.add_batch(
                    predictions=UpperCamelCase_ ,       references=UpperCamelCase_ ,       )

            A__		 =					metric.compute()
            # Use accelerator.print to print only on the main process.
            accelerator.print(f"epoch {epoch}:" ,       UpperCamelCase_	)

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
    inner_training_loop()



def 		lowerCAmelCase__    (					)->  Optional[Any]:
    A__		 =					argparse.ArgumentParser(description='''Simple example of training script.'''	)
    parser.add_argument(
        '''--mixed_precision''' ,       type=UpperCamelCase_ ,       default=UpperCamelCase_ ,       choices=['''no''', '''fp16''', '''bf16''', '''fp8'''] ,       help='''Whether to use mixed precision. Choose'''
        '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
        '''and an Nvidia Ampere GPU.''' ,       )
    parser.add_argument('''--cpu''' ,       action='''store_true''' ,       help='''If passed, will train on the CPU.'''	)
    A__		 =					parser.parse_args()
    A__		 =					{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 4_2, '''batch_size''': 1_6}
    training_function(UpperCamelCase_ ,       UpperCamelCase_	)


if __name__ == "__main__":
   main()
