






'''simple docstring'''
import json
import os
import shutil
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoConfig, BertConfig, GPTaConfig
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402


UpperCamelCase__			=      {
    'return_dict': False,
    'output_hidden_states': True,
    'output_attentions': True,
    'torchscript': True,
    'torch_dtype': 'float16',
    'use_bfloat16': True,
    'tf_legacy_loss': True,
    'pruned_heads': {'a': 1},
    'tie_word_embeddings': False,
    'is_decoder': True,
    'cross_attention_hidden_size': 128,
    'add_cross_attention': True,
    'tie_encoder_decoder': True,
    'max_length': 50,
    'min_length': 3,
    'do_sample': True,
    'early_stopping': True,
    'num_beams': 3,
    'num_beam_groups': 3,
    'diversity_penalty': 0.5,
    'temperature': 2.0,
    'top_k': 10,
    'top_p': 0.7,
    'typical_p': 0.2,
    'repetition_penalty': 0.8,
    'length_penalty': 0.8,
    'no_repeat_ngram_size': 5,
    'encoder_no_repeat_ngram_size': 5,
    'bad_words_ids': [1, 2, 3],
    'num_return_sequences': 3,
    'chunk_size_feed_forward': 5,
    'output_scores': True,
    'return_dict_in_generate': True,
    'forced_bos_token_id': 2,
    'forced_eos_token_id': 3,
    'remove_invalid_values': True,
    'architectures': ['BertModel'],
    'finetuning_task': 'translation',
    'id2label': {0: 'label'},
    'label2id': {'label': '0'},
    'tokenizer_class': 'BertTokenizerFast',
    'prefix': 'prefix',
    'bos_token_id': 6,
    'pad_token_id': 7,
    'eos_token_id': 8,
    'sep_token_id': 9,
    'decoder_start_token_id': 10,
    'exponential_decay_length_penalty': (5, 1.01),
    'suppress_tokens': [0, 1],
    'begin_suppress_tokens': 2,
    'task_specific_params': {'translation': 'some_params'},
    'problem_type': 'regression',
}


@is_staging_test
class 							_UpperCAmelCase    (		unittest.TestCase      ):




    @classmethod
    def 					lowerCAmelCase__					(      cls	:							Optional[Any]      ):






       '''simple docstring'''






       lowercase_						:     List[str]							     =  TOKEN
       HfFolder.save_token(a      )




    @classmethod
    def 					lowerCAmelCase__					(      cls	:							Dict      ):






       '''simple docstring'''






       try:
          delete_repo(token=cls._token	,  repo_id="test-config"      )
       except HTTPError:
          pass

       try:
          delete_repo(token=cls._token	,  repo_id="valid_org/test-config-org"      )
       except HTTPError:
          pass

       try:
          delete_repo(token=cls._token	,  repo_id="test-dynamic-config"      )
       except HTTPError:
          pass




    def 					lowerCAmelCase__					(      self	:							str      ):






       '''simple docstring'''






       lowercase_						:     Union[str, Any]							     =  BertConfig(
           vocab_size=9_9	,  hidden_size=3_2	,  num_hidden_layers=5	,  num_attention_heads=4	,  intermediate_size=3_7      )
       config.push_to_hub("test-config"	,  use_auth_token=self._token      )

       lowercase_						:     str							     =  BertConfig.from_pretrained(f"""{USER}/test-config"""      )
       for k, v in config.to_dict().items():
          if k != "transformers_version":
             self.assertEqual(a	,  getattr(a	,  a      )      )

        # Reset repo
       delete_repo(token=self._token	,  repo_id="test-config"      )

       # Push to hub via save_pretrained
       with tempfile.TemporaryDirectory() as tmp_dir:
          config.save_pretrained(a	,  repo_id="test-config"	,  push_to_hub=a	,  use_auth_token=self._token      )

       lowercase_						:     Any							     =  BertConfig.from_pretrained(f"""{USER}/test-config"""      )
       for k, v in config.to_dict().items():
          if k != "transformers_version":
             self.assertEqual(a	,  getattr(a	,  a      )      )




    def 					lowerCAmelCase__					(      self	:							str      ):






       '''simple docstring'''






       lowercase_						:     Tuple							     =  BertConfig(
           vocab_size=9_9	,  hidden_size=3_2	,  num_hidden_layers=5	,  num_attention_heads=4	,  intermediate_size=3_7      )
       config.push_to_hub("valid_org/test-config-org"	,  use_auth_token=self._token      )

       lowercase_						:     Dict							     =  BertConfig.from_pretrained("valid_org/test-config-org"      )
       for k, v in config.to_dict().items():
          if k != "transformers_version":
             self.assertEqual(a	,  getattr(a	,  a      )      )

        # Reset repo
       delete_repo(token=self._token	,  repo_id="valid_org/test-config-org"      )

       # Push to hub via save_pretrained
       with tempfile.TemporaryDirectory() as tmp_dir:
          config.save_pretrained(
              a	,  repo_id="valid_org/test-config-org"	,  push_to_hub=a	,  use_auth_token=self._token      )

       lowercase_						:     str							     =  BertConfig.from_pretrained("valid_org/test-config-org"      )
       for k, v in config.to_dict().items():
          if k != "transformers_version":
             self.assertEqual(a	,  getattr(a	,  a      )      )





    def 					lowerCAmelCase__					(      self	:							Optional[Any]      ):






       '''simple docstring'''






       CustomConfig.register_for_auto_class()
       lowercase_						:     List[Any]							     =  CustomConfig(attribute=4_2      )

       config.push_to_hub("test-dynamic-config"	,  use_auth_token=self._token      )

       # This has added the proper auto_map field to the config
       self.assertDictEqual(config.auto_map	,  {"AutoConfig": "custom_configuration.CustomConfig"}      )

       lowercase_						:     str							     =  AutoConfig.from_pretrained(f"""{USER}/test-dynamic-config"""	,  trust_remote_code=a      )
       # Can't make an isinstance check because the new_config is from the FakeConfig class of a dynamic module
       self.assertEqual(new_config.__class__.__name__	,  "CustomConfig"      )
       self.assertEqual(new_config.attribute	,  4_2      )


class 							_UpperCAmelCase    (		unittest.TestCase      ):




    def 					lowerCAmelCase__					(      self	:							List[str]      ):






       '''simple docstring'''






       lowercase_						:     Tuple							     =  GPTaConfig()

       # attempt to modify each of int/float/bool/str config records and verify they were updated
       lowercase_						:     Any							     =  c.n_embd + 1  # int
       lowercase_						:     List[Any]							     =  c.resid_pdrop + 1.0  # float
       lowercase_						:     Any							     =  not c.scale_attn_weights  # bool
       lowercase_						:     str							     =  c.summary_type + "foo"  # str
       c.update_from_string(
           f"""n_embd={n_embd},resid_pdrop={resid_pdrop},scale_attn_weights={scale_attn_weights},summary_type={summary_type}"""      )
       self.assertEqual(a	,  c.n_embd	,  "mismatch for key: n_embd"      )
       self.assertEqual(a	,  c.resid_pdrop	,  "mismatch for key: resid_pdrop"      )
       self.assertEqual(a	,  c.scale_attn_weights	,  "mismatch for key: scale_attn_weights"      )
       self.assertEqual(a	,  c.summary_type	,  "mismatch for key: summary_type"      )




    def 					lowerCAmelCase__					(      self	:							str      ):






       '''simple docstring'''






       lowercase_						:     str							     =  PretrainedConfig()
       lowercase_						:     Any							     =  [key for key in base_config.__dict__ if key not in config_common_kwargs]
       # If this part of the test fails, you have arguments to addin config_common_kwargs above.
       self.assertListEqual(
           a	,  ["is_encoder_decoder", "_name_or_path", "_commit_hash", "transformers_version"]      )
       lowercase_						:     Union[str, Any]							     =  [key for key, value in config_common_kwargs.items() if value == getattr(a	,  a      )]
       if len(a      ) > 0:
          raise ValueError(
              "The following keys are set with the default values in"
              " `test_configuration_common.config_common_kwargs` pick another value for them:"
              f""" {", ".join(a      )}."""      )




    def 					lowerCAmelCase__					(      self	:							Optional[Any]      ):






       '''simple docstring'''






       with self.assertRaises(a      ):
          # config is in subfolder, the following should not work without specifying the subfolder
          lowercase_						:     Dict							     =  BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert-subfolder"      )

       lowercase_						:     str							     =  BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert-subfolder"	,  subfolder="bert"      )

       self.assertIsNotNone(a      )




    def 					lowerCAmelCase__					(      self	:							Optional[Any]      ):






       '''simple docstring'''






       lowercase_						:     Tuple							     =  mock.Mock()
       lowercase_						:     Dict							     =  5_0_0
       lowercase_						:     int							     =  {}
       lowercase_						:     Tuple							     =  HTTPError
       lowercase_						:     List[str]							     =  {}

       # Download this model to make sure it's in the cache.
       lowercase_						:     Optional[int]							     =  BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert"      )

       # Under the mock environment we get a 500 error when trying to reach the model.
       with mock.patch("requests.Session.request"	,  return_value=a      ) as mock_head:
          lowercase_						:     Tuple							     =  BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert"      )
          # This check we did call the fake head request
          mock_head.assert_called()




    def 					lowerCAmelCase__					(      self	:							Optional[int]      ):






       '''simple docstring'''






       lowercase_						:     List[str]							     =  BertConfig.from_pretrained(
           "https://huggingface.co/hf-internal-testing/tiny-random-bert/resolve/main/config.json"      )




    def 					lowerCAmelCase__					(      self	:							Any      ):






       '''simple docstring'''






       lowercase_						:     Tuple							     =  AutoConfig.from_pretrained("bert-base-cased"      )
       lowercase_						:     int							     =  ["config.4.0.0.json"]

       with tempfile.TemporaryDirectory() as tmp_dir:
          configuration.save_pretrained(a      )
          lowercase_						:     List[str]							     =  2
          json.dump(configuration.to_dict()	,  open(os.path.join(a	,  "config.4.0.0.json"      )	,  "w"      )      )

          # This should pick the new configuration file as the version of Transformers is > 4.0.0
          lowercase_						:     int							     =  AutoConfig.from_pretrained(a      )
          self.assertEqual(new_configuration.hidden_size	,  2      )

          # Will need to be adjusted if we reach v42 and this test is still here.
          # Should pick the old configuration file as the version of Transformers is < 4.42.0
          lowercase_						:     List[Any]							     =  ["config.42.0.0.json"]
          lowercase_						:     Optional[int]							     =  7_6_8
          configuration.save_pretrained(a      )
          shutil.move(os.path.join(a	,  "config.4.0.0.json"      )	,  os.path.join(a	,  "config.42.0.0.json"      )      )
          lowercase_						:     Tuple							     =  AutoConfig.from_pretrained(a      )
          self.assertEqual(new_configuration.hidden_size	,  7_6_8      )





    def 					lowerCAmelCase__					(      self	:							Union[str, Any]      ):






       '''simple docstring'''






       lowercase_						:     Tuple							     =  "hf-internal-testing/test-two-configs"

       import transformers as new_transformers

       lowercase_						:     Union[str, Any]							     =  "v4.0.0"
       lowercase_     ,		lowercase_						:     int							     =  new_transformers.models.auto.AutoConfig.from_pretrained(
           a	,  return_unused_kwargs=a      )
       self.assertEqual(new_configuration.hidden_size	,  2      )
       # This checks `_configuration_file` ia not kept in the kwargs by mistake.
       self.assertDictEqual(a	,  {}      )

       # Testing an older version by monkey-patching the version in the module it's used.
       import transformers as old_transformers

       lowercase_						:     int							     =  "v3.0.0"
       lowercase_						:     Dict							     =  old_transformers.models.auto.AutoConfig.from_pretrained(a      )
       self.assertEqual(old_configuration.hidden_size	,  7_6_8      )












'''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, 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
#
########################################################################


UpperCamelCase__			=      16
UpperCamelCase__			=      32




def 		__SCREAMING_SNAKE_CASE       (	_UpperCamelCase    ,		_UpperCamelCase = 16				):


   """simple docstring"""




   lowercase_						:     int							     =  AutoTokenizer.from_pretrained("bert-base-cased"				)
   lowercase_						:     Optional[Any]							     =  load_dataset("glue"    ,		"mrpc"				)

   def tokenize_function(_UpperCamelCase				):
      # max_length=None => use the model max length (it's actually the default)
      lowercase_						:     Any							     =  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():
      lowercase_						:     str							     =  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
   lowercase_						:     str							     =  tokenized_datasets.rename_column("label"    ,		"labels"				)

   def collate_fn(_UpperCamelCase				):
      # On TPU it's best to pad everything to the same length or training will be very slow.
      lowercase_						:     List[str]							     =  128 if accelerator.distributed_type == DistributedType.TPU else None
      # When using mixed precision we want round multiples of 8/16
      if accelerator.mixed_precision == "fp8":
         lowercase_						:     List[str]							     =  16
      elif accelerator.mixed_precision != "no":
         lowercase_						:     Tuple							     =  8
      else:
         lowercase_						:     Dict							     =  None

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

   # Instantiate dataloaders.
   lowercase_						:     Optional[int]							     =  DataLoader(
       tokenized_datasets["train"]    ,		shuffle=_UpperCamelCase    ,		collate_fn=_UpperCamelCase    ,		batch_size=_UpperCamelCase				)
   lowercase_						:     str							     =  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

 UpperCamelCase__			=      mocked_dataloaders  # noqa: F811




def 		__SCREAMING_SNAKE_CASE       (	_UpperCamelCase    ,		_UpperCamelCase				):


   """simple docstring"""




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

   lowercase_						:     List[Any]							     =  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				):
      # 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)
      lowercase_						:     List[str]							     =  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).
      lowercase_						:     List[Any]							     =  model.to(accelerator.device				)

      # Instantiate optimizer
      lowercase_						:     Tuple							     =  AdamW(params=model.parameters()    ,		lr=_UpperCamelCase				)
      lowercase_     ,		lowercase_						:     Any							     =  get_dataloaders(_UpperCamelCase    ,		_UpperCamelCase				)

      # Instantiate scheduler
      lowercase_						:     Union[str, Any]							     =  get_linear_schedule_with_warmup(
          optimizer=_UpperCamelCase    ,		num_warmup_steps=100    ,		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.
      lowercase_     ,		lowercase_     ,		lowercase_     ,		lowercase_     ,		lowercase_						:     Optional[int]							     =  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				)
            lowercase_						:     Optional[Any]							     =  model(**_UpperCamelCase				)
            lowercase_						:     List[Any]							     =  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():
               lowercase_						:     Optional[int]							     =  model(**_UpperCamelCase				)
            lowercase_						:     Union[str, Any]							     =  outputs.logits.argmax(dim=-1				)
            lowercase_     ,		lowercase_						:     str							     =  accelerator.gather_for_metrics((predictions, batch["labels"])				)
            metric.add_batch(
                predictions=_UpperCamelCase    ,		references=_UpperCamelCase    ,		)

         lowercase_						:     int							     =  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 		__SCREAMING_SNAKE_CASE       (	):


   """simple docstring"""




   lowercase_						:     Optional[Any]							     =  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."				)
   lowercase_						:     int							     =  parser.parse_args()
   lowercase_						:     Optional[Any]							     =  {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
   training_function(_UpperCamelCase    ,		_UpperCamelCase				)


if __name__ == "__main__":
 main()




