





"""simple docstring"""





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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to combine both the gradient accumulation
# and automatic batch size finder utilities of Accelerate to perfrom
# automatic gradient accumulation
#
# 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
#
########################################################################

SCREAMING_SNAKE_CASE__ =       32


def       UpperCAmelCase__							(  SCREAMING_SNAKE_CASE					:						Accelerator							,						SCREAMING_SNAKE_CASE					:						int = 16      ):







      '''simple docstring'''




      lowerCAmelCase  						=	AutoTokenizer.from_pretrained("""bert-base-cased"""      )
      lowerCAmelCase  						=	load_dataset("""glue"""							,						"""mrpc"""      )

      def tokenize_function(SCREAMING_SNAKE_CASE					:						str      ):
            # max_length=None => use the model max length (it's actually the default)
            lowerCAmelCase  						=	tokenizer(examples["""sentence1"""]							,						examples["""sentence2"""]							,						truncation=_lowerCamelCase							,						max_length=_lowerCamelCase      )
            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():
            lowerCAmelCase  						=	datasets.map(
                _lowerCamelCase							,						batched=_lowerCamelCase							,						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
      lowerCAmelCase  						=	tokenized_datasets.rename_column("""label"""							,						"""labels"""      )

      def collate_fn(SCREAMING_SNAKE_CASE					:						str      ):
            # When using mixed precision we want round multiples of 8/16
            if accelerator.mixed_precision == "fp8":
                  lowerCAmelCase  						=	16
            elif accelerator.mixed_precision != "no":
                  lowerCAmelCase  						=	8
            else:
                  lowerCAmelCase  						=	None

            return tokenizer.pad(
                _lowerCamelCase							,						padding="""longest"""							,						pad_to_multiple_of=_lowerCamelCase							,						return_tensors="""pt"""							,						)

      # Instantiate dataloaders.
      lowerCAmelCase  						=	DataLoader(
          tokenized_datasets["""train"""]							,						shuffle=_lowerCamelCase							,						collate_fn=_lowerCamelCase							,						batch_size=_lowerCamelCase      )
      lowerCAmelCase  						=	DataLoader(
          tokenized_datasets["""validation"""]							,						shuffle=_lowerCamelCase							,						collate_fn=_lowerCamelCase							,						batch_size=_lowerCamelCase      )

      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

       SCREAMING_SNAKE_CASE__ =       mocked_dataloaders  # noqa: F811


def       UpperCAmelCase__							(  SCREAMING_SNAKE_CASE					:						Tuple							,						SCREAMING_SNAKE_CASE					:						str      ):







      '''simple docstring'''




      if os.environ.get("""TESTING_MOCKED_DATALOADERS"""							,						_lowerCamelCase      ) == "1":
            lowerCAmelCase  						=	2
      # Initialize accelerator
      lowerCAmelCase  						=	Accelerator(cpu=args.cpu							,						mixed_precision=args.mixed_precision      )
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      lowerCAmelCase  						=	config["""lr"""]
      lowerCAmelCase  						=	int(config["""num_epochs"""]      )
      lowerCAmelCase  						=	int(config["""seed"""]      )
      lowerCAmelCase  						=	int(config["""batch_size"""]      )

      lowerCAmelCase  						=	evaluate.load("""glue"""							,						"""mrpc"""      )

      # New Code #
      # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
      # to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
      # half each time. From this, we can calculate the number of gradient accumulation steps needed
      # and modify the Accelerator object as a result
      @find_executable_batch_size(starting_batch_size=int(_lowerCamelCase      )      )
      def inner_training_loop(SCREAMING_SNAKE_CASE					:						Optional[int]      ):
            # Since we need to modify the outside accelerator object, we need to bring it
            # to the local scope
            nonlocal accelerator

            # We can calculate the number of gradient accumulation steps based on the current
            # batch size vs the starting batch size
            lowerCAmelCase  						=	observed_batch_size // batch_size

            # And then set it in the Accelerator directly:
            lowerCAmelCase  						=	num_gradient_accumulation_steps

            # Next we need to free all of the stored model references in the Accelerator each time
            accelerator.free_memory()

            # And set the seed so our results are reproducable each reset
            set_seed(_lowerCamelCase      )

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

            # 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).
            lowerCAmelCase  						=	model.to(accelerator.device      )

            # Instantiate optimizer
            lowerCAmelCase  						=	AdamW(params=model.parameters()							,						lr=_lowerCamelCase      )
            lowerCAmelCase  						=	get_dataloaders(_lowerCamelCase							,						_lowerCamelCase      )

            # Instantiate scheduler
            lowerCAmelCase  						=	get_linear_schedule_with_warmup(
                optimizer=_lowerCamelCase							,						num_warmup_steps=1_00							,						num_training_steps=(len(_lowerCamelCase      ) * 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.
            lowerCAmelCase  						=	accelerator.prepare(
                _lowerCamelCase							,						_lowerCamelCase							,						_lowerCamelCase							,						_lowerCamelCase							,						_lowerCamelCase      )

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

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

                  lowerCAmelCase  						=	metric.compute()
                  # Use accelerator.print to print only on the main process.
                  accelerator.print(F'epoch {epoch}:'							,						_lowerCamelCase      )

            # 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       UpperCAmelCase__							(  ):







      '''simple docstring'''




      lowerCAmelCase  						=	argparse.ArgumentParser(description="""Simple example of training script."""      )
      parser.add_argument(
          """--mixed_precision"""							,						type=_lowerCamelCase							,						default=_lowerCamelCase							,						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."""      )
      lowerCAmelCase  						=	parser.parse_args()
      # New Code #
      # We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
      lowerCAmelCase  						=	{"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 2_56}
      training_function(_lowerCamelCase							,						_lowerCamelCase      )


if __name__ == "__main__":
       main()













"""simple docstring"""





from __future__ import annotations

import unittest
import warnings

from transformers import AutoTokenizer, MarianConfig, MarianTokenizer, TranslationPipeline, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

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


if is_tf_available():
       import tensorflow as tf

       from transformers import TFAutoModelForSeqaSeqLM, TFMarianModel, TFMarianMTModel

@require_tf
class 		lowercase :
      _SCREAMING_SNAKE_CASE         =				MarianConfig
      _SCREAMING_SNAKE_CASE         =				{}
      _SCREAMING_SNAKE_CASE         =				'gelu'





      def __init__(	self				,						lowercase				,						lowercase=13				,						lowercase=7				,						lowercase=True				,						lowercase=False				,						lowercase=99				,						lowercase=32				,						lowercase=2				,						lowercase=4				,						lowercase=37				,						lowercase=0.1				,						lowercase=0.1				,						lowercase=20				,						lowercase=2				,						lowercase=1				,						lowercase=0				,						)		->					int:
            lowerCAmelCase  						=	parent
            lowerCAmelCase  						=	batch_size
            lowerCAmelCase  						=	seq_length
            lowerCAmelCase  						=	is_training
            lowerCAmelCase  						=	use_labels
            lowerCAmelCase  						=	vocab_size
            lowerCAmelCase  						=	hidden_size
            lowerCAmelCase  						=	num_hidden_layers
            lowerCAmelCase  						=	num_attention_heads
            lowerCAmelCase  						=	intermediate_size

            lowerCAmelCase  						=	hidden_dropout_prob
            lowerCAmelCase  						=	attention_probs_dropout_prob
            lowerCAmelCase  						=	max_position_embeddings
            lowerCAmelCase  						=	eos_token_id
            lowerCAmelCase  						=	pad_token_id
            lowerCAmelCase  						=	bos_token_id





      def    _snake_case      (	self			)		->					List[str]:
            lowerCAmelCase  						=	ids_tensor([self.batch_size, self.seq_length - 1]				,						self.vocab_size			)
            lowerCAmelCase  						=	tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size			)				,						1			)
            lowerCAmelCase  						=	tf.concat([input_ids, eos_tensor]				,						axis=1			)

            lowerCAmelCase  						=	ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size			)

            lowerCAmelCase  						=	self.config_cls(
                vocab_size=self.vocab_size				,						d_model=self.hidden_size				,						encoder_layers=self.num_hidden_layers				,						decoder_layers=self.num_hidden_layers				,						encoder_attention_heads=self.num_attention_heads				,						decoder_attention_heads=self.num_attention_heads				,						encoder_ffn_dim=self.intermediate_size				,						decoder_ffn_dim=self.intermediate_size				,						dropout=self.hidden_dropout_prob				,						attention_dropout=self.attention_probs_dropout_prob				,						max_position_embeddings=self.max_position_embeddings				,						eos_token_ids=[2]				,						bos_token_id=self.bos_token_id				,						pad_token_id=self.pad_token_id				,						decoder_start_token_id=self.pad_token_id				,						**self.config_updates				,						)
            lowerCAmelCase  						=	prepare_marian_inputs_dict(lowercase				,						lowercase				,						lowercase			)
            return config, inputs_dict





      def    _snake_case      (	self				,						lowercase				,						lowercase			)		->					Union[str, Any]:
            lowerCAmelCase  						=	TFMarianModel(config=lowercase			).get_decoder()
            lowerCAmelCase  						=	inputs_dict["""input_ids"""]

            lowerCAmelCase  						=	input_ids[:1, :]
            lowerCAmelCase  						=	inputs_dict["""attention_mask"""][:1, :]
            lowerCAmelCase  						=	inputs_dict["""head_mask"""]
            lowerCAmelCase  						=	1

            # first forward pass
            lowerCAmelCase  						=	model(lowercase				,						attention_mask=lowercase				,						head_mask=lowercase				,						use_cache=lowercase			)

            lowerCAmelCase						,		lowerCAmelCase  						=	outputs.to_tuple()

            # create hypothetical next token and extent to next_input_ids
            lowerCAmelCase  						=	ids_tensor((self.batch_size, 3)				,						config.vocab_size			)
            lowerCAmelCase  						=	tf.cast(ids_tensor((self.batch_size, 3)				,						2			)				,						tf.inta			)

            # append to next input_ids and
            lowerCAmelCase  						=	tf.concat([input_ids, next_tokens]				,						axis=-1			)
            lowerCAmelCase  						=	tf.concat([attention_mask, next_attn_mask]				,						axis=-1			)

            lowerCAmelCase  						=	model(lowercase				,						attention_mask=lowercase			)[0]
            lowerCAmelCase  						=	model(lowercase				,						attention_mask=lowercase				,						past_key_values=lowercase			)[0]

            self.parent.assertEqual(next_tokens.shape[1]				,						output_from_past.shape[1]			)

            # select random slice
            lowerCAmelCase  						=	int(ids_tensor((1,)				,						output_from_past.shape[-1]			)			)
            lowerCAmelCase  						=	output_from_no_past[:, -3:, random_slice_idx]
            lowerCAmelCase  						=	output_from_past[:, :, random_slice_idx]

            # test that outputs are equal for slice
            tf.debugging.assert_near(lowercase				,						lowercase				,						rtol=1e-3			)







def       UpperCAmelCase__							(  SCREAMING_SNAKE_CASE					:						int							,						SCREAMING_SNAKE_CASE					:						List[str]							,						SCREAMING_SNAKE_CASE					:						Union[str, Any]							,						SCREAMING_SNAKE_CASE					:						Any=None							,						SCREAMING_SNAKE_CASE					:						str=None							,						SCREAMING_SNAKE_CASE					:						List[str]=None							,						SCREAMING_SNAKE_CASE					:						Tuple=None							,						SCREAMING_SNAKE_CASE					:						Dict=None							,						):







      '''simple docstring'''




      if attention_mask is None:
            lowerCAmelCase  						=	tf.cast(tf.math.not_equal(SCREAMING_SNAKE_CASE							,						config.pad_token_id      )							,						tf.inta      )
      if decoder_attention_mask is None:
            lowerCAmelCase  						=	tf.concat(
                [
                    tf.ones(decoder_input_ids[:, :1].shape							,						dtype=tf.inta      ),
                    tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]							,						config.pad_token_id      )							,						tf.inta      ),
                ]							,						axis=-1							,						)
      if head_mask is None:
            lowerCAmelCase  						=	tf.ones((config.encoder_layers, config.encoder_attention_heads)      )
      if decoder_head_mask is None:
            lowerCAmelCase  						=	tf.ones((config.decoder_layers, config.decoder_attention_heads)      )
      if cross_attn_head_mask is None:
            lowerCAmelCase  						=	tf.ones((config.decoder_layers, config.decoder_attention_heads)      )
      return {
          "input_ids": input_ids,
          "decoder_input_ids": decoder_input_ids,
          "attention_mask": attention_mask,
          "decoder_attention_mask": decoder_attention_mask,
          "head_mask": head_mask,
          "decoder_head_mask": decoder_head_mask,
          "cross_attn_head_mask": cross_attn_head_mask,
      }

@require_tf
class 		lowercase ( _UpperCAmelCase					,  _UpperCAmelCase					,  unittest.TestCase ):
      _SCREAMING_SNAKE_CASE         =				(TFMarianMTModel, TFMarianModel) if is_tf_available() else ()
      _SCREAMING_SNAKE_CASE         =				(TFMarianMTModel,) if is_tf_available() else ()
      _SCREAMING_SNAKE_CASE         =				(
          {
              'conversational': TFMarianMTModel,
              'feature-extraction': TFMarianModel,
              'summarization': TFMarianMTModel,
              'text2text-generation': TFMarianMTModel,
              'translation': TFMarianMTModel,
          }
          if is_tf_available()
          else {}
      )
      _SCREAMING_SNAKE_CASE         =				True
      _SCREAMING_SNAKE_CASE         =				False
      _SCREAMING_SNAKE_CASE         =				False





      def    _snake_case      (	self			)		->					Optional[int]:
            lowerCAmelCase  						=	TFMarianModelTester(self			)
            lowerCAmelCase  						=	ConfigTester(self				,						config_class=lowercase			)





      def    _snake_case      (	self			)		->					str:
            self.config_tester.run_common_tests()





      def    _snake_case      (	self			)		->					str:
            lowerCAmelCase  						=	self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.check_decoder_model_past_large_inputs(*lowercase			)







@require_tf
class 		lowercase ( unittest.TestCase ):
      _SCREAMING_SNAKE_CASE         =				1_000  # show more chars for failing integration tests





      @classmethod
      def    _snake_case      (	cls			)		->					None:
            lowerCAmelCase  						=	f'Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}'
            return cls





      @cached_property
      def    _snake_case      (	self			)		->					MarianTokenizer:
            return AutoTokenizer.from_pretrained(self.model_name			)





      @property
      def    _snake_case      (	self			)		->					int:
            return self.tokenizer.eos_token_id





      @cached_property
      def    _snake_case      (	self			)		->					Any:
            warnings.simplefilter("""error"""			)
            lowerCAmelCase  						=	TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name			)
            assert isinstance(lowercase				,						lowercase			)
            lowerCAmelCase  						=	model.config
            self.assertListEqual(c.bad_words_ids				,						[[c.pad_token_id]]			)
            self.assertEqual(c.max_length				,						512			)
            self.assertEqual(c.decoder_start_token_id				,						c.pad_token_id			)
            return model





      def    _snake_case      (	self				,						**lowercase			)		->					Tuple:
            lowerCAmelCase  						=	self.translate_src_text(**lowercase			)
            self.assertListEqual(self.expected_text				,						lowercase			)





      def    _snake_case      (	self				,						**lowercase			)		->					Union[str, Any]:
            lowerCAmelCase  						=	self.tokenizer(self.src_text				,						**lowercase				,						padding=lowercase				,						return_tensors="""tf"""			)
            lowerCAmelCase  						=	self.model.generate(
                model_inputs.input_ids				,						attention_mask=model_inputs.attention_mask				,						num_beams=2				,						max_length=128			)
            lowerCAmelCase  						=	self.tokenizer.batch_decode(generated_ids.numpy()				,						skip_special_tokens=lowercase			)
            return generated_words







@require_sentencepiece
@require_tokenizers
@require_tf
class 		lowercase ( _UpperCAmelCase ):
      _SCREAMING_SNAKE_CASE         =				'mt'
      _SCREAMING_SNAKE_CASE         =				'en'
      _SCREAMING_SNAKE_CASE         =				['Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem.']
      _SCREAMING_SNAKE_CASE         =				['Touching gently, Jesus healed a man who was affected by the sad disease of leprosy.']





      @unittest.skip("""Skipping until #12647 is resolved."""			)
      @slow
      def    _snake_case      (	self			)		->					List[Any]:
            self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class 		lowercase ( _UpperCAmelCase ):
      _SCREAMING_SNAKE_CASE         =				'en'
      _SCREAMING_SNAKE_CASE         =				'zh'
      _SCREAMING_SNAKE_CASE         =				['My name is Wolfgang and I live in Berlin']
      _SCREAMING_SNAKE_CASE         =				['我叫沃尔夫冈 我住在柏林']





      @unittest.skip("""Skipping until #12647 is resolved."""			)
      @slow
      def    _snake_case      (	self			)		->					Dict:
            self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class 		lowercase ( _UpperCAmelCase ):
      _SCREAMING_SNAKE_CASE         =				'en'
      _SCREAMING_SNAKE_CASE         =				'ROMANCE'
      _SCREAMING_SNAKE_CASE         =				[
          '>>fr<< Don\'t spend so much time watching TV.',
          '>>pt<< Your message has been sent.',
          '>>es<< He\'s two years older than me.',
      ]
      _SCREAMING_SNAKE_CASE         =				[
          'Ne passez pas autant de temps à regarder la télé.',
          'A sua mensagem foi enviada.',
          'Es dos años más viejo que yo.',
      ]





      @unittest.skip("""Skipping until #12647 is resolved."""			)
      @slow
      def    _snake_case      (	self			)		->					List[str]:
            self._assert_generated_batch_equal_expected()





      @unittest.skip("""Skipping until #12647 is resolved."""			)
      @slow
      def    _snake_case      (	self			)		->					List[str]:
            lowerCAmelCase  						=	TranslationPipeline(self.model				,						self.tokenizer				,						framework="""tf"""			)
            lowerCAmelCase  						=	pipeline(self.src_text			)
            self.assertEqual(self.expected_text				,						[x["""translation_text"""] for x in output]			)






