






"""simple docstring"""





import argparse
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


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to properly calculate the metrics on the
# validation dataset when in a distributed system, 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)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# 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
#
########################################################################


A__					: str=	16
A__					: Dict=	32







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE = 16					)					->	Dict:
    """simple docstring"""

    UpperCamelCase__         =    AutoTokenizer.from_pretrained('bert-base-cased'					)
    UpperCamelCase__         =    load_dataset('glue' ,	'mrpc'					)

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

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

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

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

    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

  A__					: Optional[Any]=	mocked_dataloaders  # noqa: F811







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE					)					->	str:
    """simple docstring"""

    if os.environ.get('TESTING_MOCKED_DATALOADERS' ,	SCREAMING_SNAKE_CASE					) == "1":
        UpperCamelCase__         =    2
    # Initialize accelerator
    UpperCamelCase__         =    Accelerator(cpu=args.cpu ,	mixed_precision=args.mixed_precision					)
    # 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__         =    evaluate.load('glue' ,	'mrpc'					)

    # If the batch size is too big we use gradient accumulation
    UpperCamelCase__         =    1
    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
        UpperCamelCase__         =    batch_size // MAX_GPU_BATCH_SIZE
        UpperCamelCase__         =    MAX_GPU_BATCH_SIZE

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

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

    # Instantiate optimizer
    UpperCamelCase__         =    AdamW(params=model.parameters() ,	lr=SCREAMING_SNAKE_CASE					)

    # Instantiate scheduler
    UpperCamelCase__         =    get_linear_schedule_with_warmup(
        optimizer=SCREAMING_SNAKE_CASE ,	num_warmup_steps=1_00 ,	num_training_steps=(len(SCREAMING_SNAKE_CASE					) * num_epochs) // gradient_accumulation_steps ,	)

    # 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(
        SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE					)

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

        model.eval()
        UpperCamelCase__         =    0
        for step, batch in enumerate(SCREAMING_SNAKE_CASE					):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            batch.to(accelerator.device					)
            with torch.no_grad():
                UpperCamelCase__         =    model(**SCREAMING_SNAKE_CASE					)
            UpperCamelCase__         =    outputs.logits.argmax(dim=-1					)
            UpperCamelCase__       ,						UpperCamelCase__         =    accelerator.gather((predictions, batch['labels'])					)
            # New Code #
            # First we check if it's a distributed system
            if accelerator.use_distributed:
                # Then see if we're on the last batch of our eval dataloader
                if step == len(SCREAMING_SNAKE_CASE					) - 1:
                    # Last batch needs to be truncated on distributed systems as it contains additional samples
                    UpperCamelCase__         =    predictions[: len(eval_dataloader.dataset					) - samples_seen]
                    UpperCamelCase__         =    references[: len(eval_dataloader.dataset					) - samples_seen]
                else:
                    # Otherwise we add the number of samples seen
                    samples_seen += references.shape[0]
            # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
            # accelerator.gather_for_metrics((predictions, batch["labels"]))
            metric.add_batch(
                predictions=SCREAMING_SNAKE_CASE ,	references=SCREAMING_SNAKE_CASE ,	)

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







def   lowerCAmelCase_( )					->	Dict:
    """simple docstring"""

    UpperCamelCase__         =    argparse.ArgumentParser(description='Simple example of training script.'					)
    parser.add_argument(
        '--mixed_precision' ,	type=SCREAMING_SNAKE_CASE ,	default=SCREAMING_SNAKE_CASE ,	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.'					)
    UpperCamelCase__         =    parser.parse_args()
    UpperCamelCase__         =    {'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
    training_function(SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE					)


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


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

if is_tf_available():
  class    __lowerCamelCase    (  tf.keras.Model  ):

     def __init__(					self ,				snake_case_      )		->       Optional[int]:
         super().__init__()
         UpperCamelCase__         =    tokenizer
         UpperCamelCase__         =    AutoConfig.from_pretrained(snake_case_      )
         UpperCamelCase__         =    TFAutoModel.from_config(snake_case_      )







     def 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_      )		->       int:
         UpperCamelCase__         =    self.tokenizer(snake_case_      )
         UpperCamelCase__         =    self.bert(**snake_case_      )
         return out["pooler_output"]
@require_tf
@require_tensorflow_text
class    __lowerCamelCase    (  unittest.TestCase  ):

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       str:
       super().setUp()

       UpperCamelCase__         =    [
           BertTokenizer.from_pretrained(snake_case_      ) for checkpoint in (TOKENIZER_CHECKPOINTS * 2)
       ]  # repeat for when fast_bert_tokenizer=false
       UpperCamelCase__         =    [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__         =    [
           '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__         =    list(zip(self.test_sentences ,				self.test_sentences[::-1]      )      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Any:
       for tokenizer, tf_tokenizer in zip(self.tokenizers ,				self.tf_tokenizers      ):
           for test_inputs in (self.test_sentences, self.paired_sentences):
               UpperCamelCase__         =    tokenizer(snake_case_ ,				return_tensors='tf' ,				padding='longest'      )
               UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self      )		->       str:
       for tf_tokenizer in self.tf_tokenizers:
           UpperCamelCase__         =    tf_tokenizer(self.paired_sentences      )
           UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self      )		->       Dict:
       for tf_tokenizer in self.tf_tokenizers:
           UpperCamelCase__         =    tf.function(snake_case_      )
           for test_inputs in (self.test_sentences, self.paired_sentences):
               UpperCamelCase__         =    tf.constant(snake_case_      )
               UpperCamelCase__         =    compiled_tokenizer(snake_case_      )
               UpperCamelCase__         =    tf_tokenizer(snake_case_      )

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







   @slow
   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       List[str]:
       for tf_tokenizer in self.tf_tokenizers:
           UpperCamelCase__         =    ModelToSave(tokenizer=snake_case_      )
           UpperCamelCase__         =    tf.convert_to_tensor(self.test_sentences      )
           UpperCamelCase__         =    model(snake_case_      )  # Build model with some sample inputs
           with TemporaryDirectory() as tempdir:
               UpperCamelCase__         =    Path(snake_case_      ) / 'saved.model'
               model.save(snake_case_      )
               UpperCamelCase__         =    tf.keras.models.load_model(snake_case_      )
           UpperCamelCase__         =    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      )




