



"""simple docstring"""


import argparse
import json
import os

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.deepspeed import DummyOptim, DummyScheduler


snake_case_    					=	16
snake_case_    					=	32

def 			_lowerCAmelCase				(       lowercase_		,     lowercase_ = 16		,     lowercase_ = "bert-base-cased"						):
      UpperCAmelCase					=			AutoTokenizer.from_pretrained(__lowerCamelCase						)
      UpperCAmelCase					=			load_dataset('glue'		,     'mrpc'						)

      def tokenize_function(lowercase_						):
            # max_length=None => use the model max length (it's actually the default)
            UpperCAmelCase					=			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
      UpperCAmelCase					=			datasets.map(
          __lowerCamelCase		,     batched=__lowerCamelCase		,     remove_columns=['idx', 'sentence1', 'sentence2']		,     load_from_cache_file=__lowerCamelCase						)

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

      def collate_fn(lowercase_						):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            if accelerator.distributed_type == DistributedType.TPU:
                  return tokenizer.pad(__lowerCamelCase		,     padding='max_length'		,     max_length=128		,     return_tensors='pt'						)
            return tokenizer.pad(__lowerCamelCase		,     padding='longest'		,     return_tensors='pt'						)

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

      return train_dataloader, eval_dataloader

def 			_lowerCAmelCase				(       lowercase_		,     lowercase_		,     lowercase_		,     lowercase_						):
      model.eval()
      UpperCAmelCase					=			0
      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():
                  UpperCAmelCase					=			model(**__lowerCamelCase						)
            UpperCAmelCase					=			outputs.logits.argmax(dim=-1						)
            # It is slightly faster to call this once, than multiple times
            UpperCAmelCase				,   UpperCAmelCase					=			accelerator.gather(
                (predictions, batch['labels'])						)  # If we are in a multiprocess environment, the last batch has duplicates
            if accelerator.use_distributed:
                  if step == len(__lowerCamelCase						) - 1:
                        UpperCAmelCase					=			predictions[: len(eval_dataloader.dataset						) - samples_seen]
                        UpperCAmelCase					=			references[: len(eval_dataloader.dataset						) - samples_seen]
                  else:
                        samples_seen += references.shape[0]
            metric.add_batch(
                predictions=__lowerCamelCase		,     references=__lowerCamelCase		,     )

      UpperCAmelCase					=			metric.compute()
      return eval_metric["accuracy"]

def 			_lowerCAmelCase				(       lowercase_		,     lowercase_						):
      # Initialize accelerator
      UpperCAmelCase					=			Accelerator()

      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      UpperCAmelCase					=			config['lr']
      UpperCAmelCase					=			int(config['num_epochs']						)
      UpperCAmelCase					=			int(config['seed']						)
      UpperCAmelCase					=			int(config['batch_size']						)
      UpperCAmelCase					=			args.model_name_or_path

      set_seed(__lowerCamelCase						)
      UpperCAmelCase				,   UpperCAmelCase					=			get_dataloaders(__lowerCamelCase		,     __lowerCamelCase		,     __lowerCamelCase						)

      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      UpperCAmelCase					=			AutoModelForSequenceClassification.from_pretrained(__lowerCamelCase		,     return_dict=__lowerCamelCase						)

      # Instantiate optimizer
      UpperCAmelCase					=			(
          AdamW
          if accelerator.state.deepspeed_plugin is None
          or 'optimizer' not in accelerator.state.deepspeed_plugin.deepspeed_config
          else DummyOptim
      )
      UpperCAmelCase					=			optimizer_cls(params=model.parameters()		,     lr=__lowerCamelCase						)

      if accelerator.state.deepspeed_plugin is not None:
            UpperCAmelCase					=			accelerator.state.deepspeed_plugin.deepspeed_config[
                'gradient_accumulation_steps'
            ]
      else:
            UpperCAmelCase					=			1
      UpperCAmelCase					=			(len(__lowerCamelCase						) * num_epochs) // gradient_accumulation_steps

      # Instantiate scheduler
      if (
          accelerator.state.deepspeed_plugin is None
          or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
      ):
            UpperCAmelCase					=			get_linear_schedule_with_warmup(
                optimizer=__lowerCamelCase		,     num_warmup_steps=0		,     num_training_steps=__lowerCamelCase		,     )
      else:
            UpperCAmelCase					=			DummyScheduler(__lowerCamelCase		,     total_num_steps=__lowerCamelCase		,     warmup_num_steps=0						)

      # 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.
      UpperCAmelCase				,   UpperCAmelCase				,   UpperCAmelCase				,   UpperCAmelCase				,   UpperCAmelCase					=			accelerator.prepare(
          __lowerCamelCase		,     __lowerCamelCase		,     __lowerCamelCase		,     __lowerCamelCase		,     __lowerCamelCase						)

      # We need to keep track of how many total steps we have iterated over
      UpperCAmelCase					=			0
      # We also need to keep track of the stating epoch so files are named properly
      UpperCAmelCase					=			0
      UpperCAmelCase					=			evaluate.load('glue'		,     'mrpc'						)
      UpperCAmelCase					=			num_epochs

      if args.partial_train_epoch is not None:
            UpperCAmelCase					=			args.partial_train_epoch

      if args.resume_from_checkpoint:
            accelerator.load_state(args.resume_from_checkpoint						)
            UpperCAmelCase					=			args.resume_from_checkpoint.split('epoch_'						)[1]
            UpperCAmelCase					=			''
            for char in epoch_string:
                  if char.isdigit():
                        state_epoch_num += char
                  else:
                        break
            UpperCAmelCase					=			int(__lowerCamelCase						) + 1
            UpperCAmelCase					=			evaluation_loop(__lowerCamelCase		,     __lowerCamelCase		,     __lowerCamelCase		,     __lowerCamelCase						)
            accelerator.print('resumed checkpoint performance:'		,     __lowerCamelCase						)
            accelerator.print('resumed checkpoint\'s scheduler\'s lr:'		,     lr_scheduler.get_lr()[0]						)
            accelerator.print('resumed optimizers\'s lr:'		,     optimizer.param_groups[0]['lr']						)
            with open(os.path.join(args.output_dir		,     F"""state_{starting_epoch-1}.json"""						)		,     'r'						) as f:
                  UpperCAmelCase					=			json.load(__lowerCamelCase						)
                  assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed"
                  assert (
                      resumed_state["lr"] == lr_scheduler.get_lr()[0]
                  ), "Scheduler learning rate mismatch, loading from checkpoint failed"
                  assert (
                      resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"]
                  ), "Optimizer learning rate mismatch, loading from checkpoint failed"
                  assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed"
                  return

            # Now we train the model
      UpperCAmelCase					=			{}
      for epoch in range(__lowerCamelCase		,     __lowerCamelCase						):
            model.train()
            for step, batch in enumerate(__lowerCamelCase						):
                  UpperCAmelCase					=			model(**__lowerCamelCase						)
                  UpperCAmelCase					=			outputs.loss
                  UpperCAmelCase					=			loss / gradient_accumulation_steps
                  accelerator.backward(__lowerCamelCase						)
                  if step % gradient_accumulation_steps == 0:
                        optimizer.step()
                        lr_scheduler.step()
                        optimizer.zero_grad()

                  overall_step += 1
            UpperCAmelCase					=			F"""epoch_{epoch}"""
            UpperCAmelCase					=			os.path.join(args.output_dir		,     __lowerCamelCase						)
            accelerator.save_state(__lowerCamelCase						)
            UpperCAmelCase					=			evaluation_loop(__lowerCamelCase		,     __lowerCamelCase		,     __lowerCamelCase		,     __lowerCamelCase						)
            UpperCAmelCase					=			accuracy
            UpperCAmelCase					=			lr_scheduler.get_lr()[0]
            UpperCAmelCase					=			optimizer.param_groups[0]['lr']
            UpperCAmelCase					=			epoch
            UpperCAmelCase					=			overall_step
            accelerator.print(F"""epoch {epoch}:"""		,     __lowerCamelCase						)

            accelerator.wait_for_everyone()
            if accelerator.is_main_process:
                  with open(os.path.join(args.output_dir		,     F"""state_{epoch}.json"""						)		,     'w'						) as f:
                        json.dump(__lowerCamelCase		,     __lowerCamelCase						)

def 			_lowerCAmelCase				(       ):
      UpperCAmelCase					=			argparse.ArgumentParser(description='Simple example of training script tracking peak GPU memory usage.'						)
      parser.add_argument(
          '--model_name_or_path'		,     type=__lowerCamelCase		,     default='bert-base-cased'		,     help='Path to pretrained model or model identifier from huggingface.co/models.'		,     required=__lowerCamelCase		,     )
      parser.add_argument(
          '--output_dir'		,     type=__lowerCamelCase		,     default='.'		,     help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'		,     )
      parser.add_argument(
          '--resume_from_checkpoint'		,     type=__lowerCamelCase		,     default=__lowerCamelCase		,     help='If the training should continue from a checkpoint folder.'		,     )
      parser.add_argument(
          '--partial_train_epoch'		,     type=__lowerCamelCase		,     default=__lowerCamelCase		,     help='If passed, the training will stop after this number of epochs.'		,     )
      parser.add_argument(
          '--num_epochs'		,     type=__lowerCamelCase		,     default=2		,     help='Number of train epochs.'		,     )
      UpperCAmelCase					=			parser.parse_args()
      UpperCAmelCase					=			{'lr': 2e-5, 'num_epochs': args.num_epochs, 'seed': 42, 'batch_size': 16}

      training_function(__lowerCamelCase		,     __lowerCamelCase						)


if __name__ == "__main__":
    main()










"""simple docstring"""


import importlib
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
import transformers.models.auto
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


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

from test_module.custom_configuration import CustomConfig  # noqa E402


snake_case_    					=	get_tests_dir("""fixtures/dummy-config.json""")







class 				A_      (					unittest.TestCase    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:int )    ->      Optional[Any]:
												UpperCAmelCase					=			0






						def      UpperCAmelCase__					( self			:List[str] )    ->      str:
												self.assertIsNotNone(transformers.models.auto.__spec__ )
												self.assertIsNotNone(importlib.util.find_spec('transformers.models.auto' ) )






						def      UpperCAmelCase__					( self			:List[Any] )    ->      List[str]:
												UpperCAmelCase					=			AutoConfig.from_pretrained('bert-base-uncased' )
												self.assertIsInstance(lowercase_       ,  lowercase_ )






						def      UpperCAmelCase__					( self			:Optional[int] )    ->      int:
												UpperCAmelCase					=			AutoConfig.from_pretrained(lowercase_ )
												self.assertIsInstance(lowercase_       ,  lowercase_ )






						def      UpperCAmelCase__					( self			:int )    ->      Any:
												UpperCAmelCase					=			AutoConfig.from_pretrained(lowercase_ )
												self.assertIsInstance(lowercase_       ,  lowercase_ )






						def      UpperCAmelCase__					( self			:Optional[Any] )    ->      List[str]:
												UpperCAmelCase					=			AutoConfig.for_model('roberta' )
												self.assertIsInstance(lowercase_       ,  lowercase_ )






						def      UpperCAmelCase__					( self			:str )    ->      List[str]:
												with tempfile.TemporaryDirectory() as tmp_dir:
																		# This model name contains bert and roberta, but roberta ends up being picked.
																		UpperCAmelCase					=			os.path.join(lowercase_       ,  'fake-roberta' )
																		os.makedirs(lowercase_       ,  exist_ok=lowercase_ )
																		with open(os.path.join(lowercase_       ,  'config.json' )       ,  'w' ) as f:
																								f.write(json.dumps({} ) )
																		UpperCAmelCase					=			AutoConfig.from_pretrained(lowercase_ )
																		self.assertEqual(type(lowercase_ )       ,  lowercase_ )






						def      UpperCAmelCase__					( self			:int )    ->      Union[str, Any]:
												try:
																		AutoConfig.register('custom'       ,  lowercase_ )
																		# Wrong model type will raise an error
																		with self.assertRaises(lowercase_ ):
																								AutoConfig.register('model'       ,  lowercase_ )
																		# Trying to register something existing in the Transformers library will raise an error
																		with self.assertRaises(lowercase_ ):
																								AutoConfig.register('bert'       ,  lowercase_ )

																		# Now that the config is registered, it can be used as any other config with the auto-API
																		UpperCAmelCase					=			CustomConfig()
																		with tempfile.TemporaryDirectory() as tmp_dir:
																								config.save_pretrained(lowercase_ )
																								UpperCAmelCase					=			AutoConfig.from_pretrained(lowercase_ )
																								self.assertIsInstance(lowercase_       ,  lowercase_ )

												finally:
																		if "custom" in CONFIG_MAPPING._extra_content:
																								del CONFIG_MAPPING._extra_content["custom"]






						def      UpperCAmelCase__					( self			:str )    ->      Dict:
												with self.assertRaisesRegex(
												    lowercase_       ,  'bert-base is not a local folder and is not a valid model identifier' ):
																		UpperCAmelCase					=			AutoConfig.from_pretrained('bert-base' )






						def      UpperCAmelCase__					( self			:List[Any] )    ->      Optional[int]:
												with self.assertRaisesRegex(
												    lowercase_       ,  R'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)' ):
																		UpperCAmelCase					=			AutoConfig.from_pretrained(lowercase_       ,  revision='aaaaaa' )






						def      UpperCAmelCase__					( self			:List[str] )    ->      str:
												with self.assertRaisesRegex(
												    lowercase_       ,  'hf-internal-testing/no-config-test-repo does not appear to have a file named config.json.'       ,  ):
																		UpperCAmelCase					=			AutoConfig.from_pretrained('hf-internal-testing/no-config-test-repo' )






						def      UpperCAmelCase__					( self			:str )    ->      int:
												# If remote code is not set, we will time out when asking whether to load the model.
												with self.assertRaises(lowercase_ ):
																		UpperCAmelCase					=			AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model' )
												# If remote code is disabled, we can't load this config.
												with self.assertRaises(lowercase_ ):
																		UpperCAmelCase					=			AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'       ,  trust_remote_code=lowercase_ )

												UpperCAmelCase					=			AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'       ,  trust_remote_code=lowercase_ )
												self.assertEqual(config.__class__.__name__       ,  'NewModelConfig' )

												# Test config can be reloaded.
												with tempfile.TemporaryDirectory() as tmp_dir:
																		config.save_pretrained(lowercase_ )
																		UpperCAmelCase					=			AutoConfig.from_pretrained(lowercase_       ,  trust_remote_code=lowercase_ )
												self.assertEqual(reloaded_config.__class__.__name__       ,  'NewModelConfig' )






						def      UpperCAmelCase__					( self			:List[Any] )    ->      Optional[int]:







												class 				A_      (					SCREAMING_SNAKE_CASE_    ):



																		"""simple docstring"""
																		__UpperCamelCase		    =							"""new-model"""

												try:
																		AutoConfig.register('new-model'       ,  lowercase_ )
																		# If remote code is not set, the default is to use local
																		UpperCAmelCase					=			AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model' )
																		self.assertEqual(config.__class__.__name__       ,  'NewModelConfigLocal' )

																		# If remote code is disabled, we load the local one.
																		UpperCAmelCase					=			AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'       ,  trust_remote_code=lowercase_ )
																		self.assertEqual(config.__class__.__name__       ,  'NewModelConfigLocal' )

																		# If remote is enabled, we load from the Hub
																		UpperCAmelCase					=			AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'       ,  trust_remote_code=lowercase_ )
																		self.assertEqual(config.__class__.__name__       ,  'NewModelConfig' )

												finally:
																		if "new-model" in CONFIG_MAPPING._extra_content:
																								del CONFIG_MAPPING._extra_content["new-model"]





