


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


__UpperCAmelCase    :Optional[Any]				       =     1_6
__UpperCAmelCase    :List[str]				       =     3_2




def 			_a			(		_lowercase      :			Accelerator    ,      _lowercase      :			int = 16    ,      _lowercase      :			str = "bert-base-cased"						):





     '''simple docstring'''






     __UpperCAmelCase			:				Union[str, Any]													=    AutoTokenizer.from_pretrained(_lowercase						)
     __UpperCAmelCase			:				List[str]													=    load_dataset('''glue'''    ,      '''mrpc'''						)

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

     # Apply the method we just defined to all the examples in all the splits of the dataset
     __UpperCAmelCase			:				List[str]													=    datasets.map(
         _lowercase    ,      batched=_lowercase    ,      remove_columns=['''idx''', '''sentence1''', '''sentence2''']    ,      load_from_cache_file=_lowercase						)

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

     def collate_fn(_lowercase      :			Any						):
          # 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(_lowercase    ,      padding='''max_length'''    ,      max_length=128    ,      return_tensors='''pt'''						)
          return tokenizer.pad(_lowercase    ,      padding='''longest'''    ,      return_tensors='''pt'''						)

     # Instantiate dataloaders.
     __UpperCAmelCase			:				int													=    DataLoader(
         tokenized_datasets['''train''']    ,      shuffle=_lowercase    ,      collate_fn=_lowercase    ,      batch_size=_lowercase						)
     __UpperCAmelCase			:				Dict													=    DataLoader(
         tokenized_datasets['''validation''']    ,      shuffle=_lowercase    ,      collate_fn=_lowercase    ,      batch_size=_lowercase						)

     return train_dataloader, eval_dataloader




def 			_a			(		_lowercase      :			str    ,      _lowercase      :			Any    ,      _lowercase      :			Any    ,      _lowercase      :			Any						):





     '''simple docstring'''






     model.eval()
     __UpperCAmelCase			:				Union[str, Any]													=    0
     for step, batch in enumerate(_lowercase						):
          # We could avoid this line since we set the accelerator with `device_placement=True`.
          batch.to(accelerator.device						)
          with torch.no_grad():
               __UpperCAmelCase			:				Dict													=    model(**_lowercase						)
          __UpperCAmelCase			:				int													=    outputs.logits.argmax(dim=-1						)
          # It is slightly faster to call this once, than multiple times
          __UpperCAmelCase							,     __UpperCAmelCase			:				Optional[int]													=    accelerator.gather(
              (predictions, batch['''labels'''])						)  # If we are in a multiprocess environment, the last batch has duplicates
          if accelerator.use_distributed:
               if step == len(_lowercase						) - 1:
                    __UpperCAmelCase			:				List[str]													=    predictions[: len(eval_dataloader.dataset						) - samples_seen]
                    __UpperCAmelCase			:				Tuple													=    references[: len(eval_dataloader.dataset						) - samples_seen]
               else:
                    samples_seen += references.shape[0]
          metric.add_batch(
              predictions=_lowercase    ,      references=_lowercase    ,      )

     __UpperCAmelCase			:				Union[str, Any]													=    metric.compute()
     return eval_metric["accuracy"]




def 			_a			(		_lowercase      :			Optional[Any]    ,      _lowercase      :			List[Any]						):





     '''simple docstring'''






     __UpperCAmelCase			:				List[str]													=    Accelerator()

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

     set_seed(_lowercase						)
     __UpperCAmelCase							,     __UpperCAmelCase			:				Optional[Any]													=    get_dataloaders(_lowercase    ,      _lowercase    ,      _lowercase						)

     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     __UpperCAmelCase			:				Dict													=    AutoModelForSequenceClassification.from_pretrained(_lowercase    ,      return_dict=_lowercase						)

     # Instantiate optimizer
     __UpperCAmelCase			:				List[Any]													=    (
         AdamW
         if accelerator.state.deepspeed_plugin is None
         or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
         else DummyOptim
     )
     __UpperCAmelCase			:				Optional[Any]													=    optimizer_cls(params=model.parameters()    ,      lr=_lowercase						)

     if accelerator.state.deepspeed_plugin is not None:
          __UpperCAmelCase			:				Union[str, Any]													=    accelerator.state.deepspeed_plugin.deepspeed_config[
              '''gradient_accumulation_steps'''
          ]
     else:
          __UpperCAmelCase			:				Union[str, Any]													=    1
     __UpperCAmelCase			:				str													=    (len(_lowercase						) * 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			:				Any													=    get_linear_schedule_with_warmup(
              optimizer=_lowercase    ,      num_warmup_steps=0    ,      num_training_steps=_lowercase    ,      )
     else:
          __UpperCAmelCase			:				Optional[Any]													=    DummyScheduler(_lowercase    ,      total_num_steps=_lowercase    ,      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			:				Optional[Any]													=    accelerator.prepare(
         _lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase						)

     # We need to keep track of how many total steps we have iterated over
     __UpperCAmelCase			:				Optional[int]													=    0
     # We also need to keep track of the stating epoch so files are named properly
     __UpperCAmelCase			:				Any													=    0
     __UpperCAmelCase			:				int													=    evaluate.load('''glue'''    ,      '''mrpc'''						)
     __UpperCAmelCase			:				List[str]													=    num_epochs

     if args.partial_train_epoch is not None:
          __UpperCAmelCase			:				List[Any]													=    args.partial_train_epoch

     if args.resume_from_checkpoint:
          accelerator.load_state(args.resume_from_checkpoint						)
          __UpperCAmelCase			:				List[Any]													=    args.resume_from_checkpoint.split('''epoch_'''						)[1]
          __UpperCAmelCase			:				str													=    ''''''
          for char in epoch_string:
               if char.isdigit():
                    state_epoch_num += char
               else:
                    break
          __UpperCAmelCase			:				List[str]													=    int(_lowercase						) + 1
          __UpperCAmelCase			:				Any													=    evaluation_loop(_lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase						)
          accelerator.print('''resumed checkpoint performance:'''    ,      _lowercase						)
          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			:				Union[str, Any]													=    json.load(_lowercase						)
               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			:				str													=    {}
     for epoch in range(_lowercase    ,      _lowercase						):
          model.train()
          for step, batch in enumerate(_lowercase						):
               __UpperCAmelCase			:				Union[str, Any]													=    model(**_lowercase						)
               __UpperCAmelCase			:				Optional[Any]													=    outputs.loss
               __UpperCAmelCase			:				Tuple													=    loss / gradient_accumulation_steps
               accelerator.backward(_lowercase						)
               if step % gradient_accumulation_steps == 0:
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

               overall_step += 1
          __UpperCAmelCase			:				Optional[int]													=    F'epoch_{epoch}'
          __UpperCAmelCase			:				Tuple													=    os.path.join(args.output_dir    ,      _lowercase						)
          accelerator.save_state(_lowercase						)
          __UpperCAmelCase			:				Any													=    evaluation_loop(_lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase						)
          __UpperCAmelCase			:				Any													=    accuracy
          __UpperCAmelCase			:				Union[str, Any]													=    lr_scheduler.get_lr()[0]
          __UpperCAmelCase			:				Optional[int]													=    optimizer.param_groups[0]['''lr''']
          __UpperCAmelCase			:				Dict													=    epoch
          __UpperCAmelCase			:				Tuple													=    overall_step
          accelerator.print(F'epoch {epoch}:'    ,      _lowercase						)

          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(_lowercase    ,      _lowercase						)




def 			_a			(		):





     '''simple docstring'''






     __UpperCAmelCase			:				Any													=    argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.'''						)
     parser.add_argument(
         '''--model_name_or_path'''    ,      type=_lowercase    ,      default='''bert-base-cased'''    ,      help='''Path to pretrained model or model identifier from huggingface.co/models.'''    ,      required=_lowercase    ,      )
     parser.add_argument(
         '''--output_dir'''    ,      type=_lowercase    ,      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=_lowercase    ,      default=_lowercase    ,      help='''If the training should continue from a checkpoint folder.'''    ,      )
     parser.add_argument(
         '''--partial_train_epoch'''    ,      type=_lowercase    ,      default=_lowercase    ,      help='''If passed, the training will stop after this number of epochs.'''    ,      )
     parser.add_argument(
         '''--num_epochs'''    ,      type=_lowercase    ,      default=2    ,      help='''Number of train epochs.'''    ,      )
     __UpperCAmelCase			:				str													=    parser.parse_args()
     __UpperCAmelCase			:				Tuple													=    {'''lr''': 2E-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}

     training_function(_lowercase    ,      _lowercase						)


if __name__ == "__main__":
      main()



'''simple docstring'''



import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFAutoModel, is_tensorflow_text_available, is_tf_available
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.testing_utils import require_tensorflow_text, require_tf, slow


if is_tf_available():
      import tensorflow as tf

if is_tensorflow_text_available():
      from transformers.models.bert import TFBertTokenizer


__UpperCAmelCase    :Optional[int]				       =     ["bert-base-uncased", "bert-base-cased"]
__UpperCAmelCase    :str				       =     "hf-internal-testing/tiny-bert-tf-only"

if is_tf_available():



      class 			a       (	tf.keras.Model ):


           """simple docstring"""






           def __init__(       self  :	List[str]							,				snake_case  :	List[str]    )		->							str:
                super().__init__()
                __UpperCAmelCase			:				List[str]													=    tokenizer
                __UpperCAmelCase			:				List[Any]													=    AutoConfig.from_pretrained(snake_case    )
                __UpperCAmelCase			:				int													=    TFAutoModel.from_config(snake_case    )






           def 				lowerCamelCase__				(       self  :	List[Any]							,				snake_case  :	Optional[int]    )		->							Optional[Any]:
                __UpperCAmelCase			:				List[Any]													=    self.tokenizer(snake_case    )
                __UpperCAmelCase			:				Optional[Any]													=    self.bert(**snake_case    )
                return out["pooler_output"]



@require_tf
@require_tensorflow_text
class 			a       (	unittest.TestCase ):


     """simple docstring"""






     def 				lowerCamelCase__				(       self  :	Optional[int]    )		->							List[str]:
          super().setUp()

          __UpperCAmelCase			:				Tuple													=    [
              BertTokenizer.from_pretrained(snake_case    ) for checkpoint in (TOKENIZER_CHECKPOINTS * 2)
          ]  # repeat for when fast_bert_tokenizer=false
          __UpperCAmelCase			:				Any													=    [TFBertTokenizer.from_pretrained(snake_case    ) for checkpoint in TOKENIZER_CHECKPOINTS] + [
              TFBertTokenizer.from_pretrained(snake_case							,				use_fast_bert_tokenizer=snake_case    )
              for checkpoint in TOKENIZER_CHECKPOINTS
          ]
          assert len(self.tokenizers    ) == len(self.tf_tokenizers    )

          __UpperCAmelCase			:				Optional[int]													=    [
              '''This is a straightforward English test sentence.''',
              '''This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.''',
              '''Now we\'re going to add some Chinese: 一 二 三 一二三''',
              '''And some much more rare Chinese: 齉 堃 齉堃''',
              '''Je vais aussi écrire en français pour tester les accents''',
              '''Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ''',
          ]
          __UpperCAmelCase			:				Optional[int]													=    list(zip(self.test_sentences							,				self.test_sentences[::-1]    )    )

     def 				lowerCamelCase__				(       self  :	Optional[int]    )		->							Optional[Any]:
          for tokenizer, tf_tokenizer in zip(self.tokenizers							,				self.tf_tokenizers    ):
               for test_inputs in (self.test_sentences, self.paired_sentences):
                    __UpperCAmelCase			:				Any													=    tokenizer(snake_case							,				return_tensors='''tf'''							,				padding='''longest'''    )
                    __UpperCAmelCase			:				Optional[int]													=    tf_tokenizer(snake_case    )

                    for key in python_outputs.keys():
                         self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape    )    )
                         self.assertTrue(tf.reduce_all(tf.cast(python_outputs[key]							,				tf.intaa    ) == tf_outputs[key]    )    )

     @slow
     def 				lowerCamelCase__				(       self  :	List[Any]    )		->							str:
          for tf_tokenizer in self.tf_tokenizers:
               __UpperCAmelCase			:				Any													=    tf_tokenizer(self.paired_sentences    )
               __UpperCAmelCase			:				Union[str, Any]													=    tf_tokenizer(
                   text=[sentence[0] for sentence in self.paired_sentences]							,				text_pair=[sentence[1] for sentence in self.paired_sentences]							,				)
               for key in merged_outputs.keys():
                    self.assertTrue(tf.reduce_all(tf.cast(merged_outputs[key]							,				tf.intaa    ) == separated_outputs[key]    )    )

     @slow
     def 				lowerCamelCase__				(       self  :	str    )		->							Union[str, Any]:
          for tf_tokenizer in self.tf_tokenizers:
               __UpperCAmelCase			:				Optional[int]													=    tf.function(snake_case    )
               for test_inputs in (self.test_sentences, self.paired_sentences):
                    __UpperCAmelCase			:				int													=    tf.constant(snake_case    )
                    __UpperCAmelCase			:				Tuple													=    compiled_tokenizer(snake_case    )
                    __UpperCAmelCase			:				Optional[int]													=    tf_tokenizer(snake_case    )

                    for key in eager_outputs.keys():
                         self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]    )    )






     @slow
     def 				lowerCamelCase__				(       self  :	str    )		->							str:
          for tf_tokenizer in self.tf_tokenizers:
               __UpperCAmelCase			:				List[Any]													=    ModelToSave(tokenizer=snake_case    )
               __UpperCAmelCase			:				Union[str, Any]													=    tf.convert_to_tensor(self.test_sentences    )
               __UpperCAmelCase			:				Tuple													=    model(snake_case    )  # Build model with some sample inputs
               with TemporaryDirectory() as tempdir:
                    __UpperCAmelCase			:				Any													=    Path(snake_case    ) / '''saved.model'''
                    model.save(snake_case    )
                    __UpperCAmelCase			:				str													=    tf.keras.models.load_model(snake_case    )
               __UpperCAmelCase			:				Optional[int]													=    loaded_model(snake_case    )
               # We may see small differences because the loaded model is compiled, so we need an epsilon for the test
               self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output    )    )							,				1E-5    )