






'''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
# and perform 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)
#
# 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_											=						16
a_											=						32

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE = 1_6       ):
    """simple docstring"""






    snake_case_ :    List[Any]     			=						AutoTokenizer.from_pretrained("bert-base-cased"       )
    snake_case_ :    List[Any]     			=						load_dataset("glue",		"mrpc"       )

    def tokenize_function(__SCREAMING_SNAKE_CASE       ):
        # max_length=None => use the model max length (it's actually the default)
        snake_case_ :    Optional[Any]     			=						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():
        snake_case_ :    List[Any]     			=						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
    snake_case_ :    Optional[int]     			=						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.
        snake_case_ :    int     			=						1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
        # When using mixed precision we want round multiples of 8/16
        if accelerator.mixed_precision == "fp8":
            snake_case_ :    str     			=						1_6
        elif accelerator.mixed_precision != "no":
            snake_case_ :    str     			=						8
        else:
            snake_case_ :    int     			=						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.
    snake_case_ :    Tuple     			=						DataLoader(
        tokenized_datasets["train"],		shuffle=__SCREAMING_SNAKE_CASE,		collate_fn=__SCREAMING_SNAKE_CASE,		batch_size=__SCREAMING_SNAKE_CASE       )
    snake_case_ :    Optional[int]     			=						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_											=						mocked_dataloaders  # noqa: F811

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    if os.environ.get("TESTING_MOCKED_DATALOADERS",		__SCREAMING_SNAKE_CASE       ) == "1":
        snake_case_ :    List[str]     			=						2
    # New Code #
    snake_case_ :    Optional[int]     			=						int(args.gradient_accumulation_steps       )
    # Initialize accelerator
    snake_case_ :    Optional[int]     			=						Accelerator(
        cpu=args.cpu,		mixed_precision=args.mixed_precision,		gradient_accumulation_steps=__SCREAMING_SNAKE_CASE       )
    if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
        raise NotImplementedError(
            "Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`"       )
    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    snake_case_ :    Optional[int]     			=						config["lr"]
    snake_case_ :    Union[str, Any]     			=						int(config["num_epochs"]       )
    snake_case_ :    Union[str, Any]     			=						int(config["seed"]       )
    snake_case_ :    Any     			=						int(config["batch_size"]       )

    snake_case_ :    int     			=						evaluate.load("glue",		"mrpc"       )

    set_seed(__SCREAMING_SNAKE_CASE       )
    snake_case_ :    str     			=						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)
    snake_case_ :    Dict     			=						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).
    snake_case_ :    Tuple     			=						model.to(accelerator.device       )

    # Instantiate optimizer
    snake_case_ :    List[str]     			=						AdamW(params=model.parameters(),		lr=__SCREAMING_SNAKE_CASE       )

    # Instantiate scheduler
    snake_case_ :    List[Any]     			=						get_linear_schedule_with_warmup(
        optimizer=__SCREAMING_SNAKE_CASE,		num_warmup_steps=1_0_0,		num_training_steps=(len(__SCREAMING_SNAKE_CASE       ) * 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.
    snake_case_ :    Union[str, Any]     			=						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       )
            # New code #
            # We use the new `accumulate` context manager to perform gradient accumulation
            # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
            with accelerator.accumulate(__SCREAMING_SNAKE_CASE       ):
                snake_case_ :    Optional[Any]     			=						model(**__SCREAMING_SNAKE_CASE       )
                snake_case_ :    str     			=						output.loss
                accelerator.backward(__SCREAMING_SNAKE_CASE       )
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

        model.eval()
        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():
                snake_case_ :    int     			=						model(**__SCREAMING_SNAKE_CASE       )
            snake_case_ :    int     			=						outputs.logits.argmax(dim=-1       )
            snake_case_ :    Dict     			=						accelerator.gather_for_metrics((predictions, batch["labels"])       )
            metric.add_batch(
                predictions=__SCREAMING_SNAKE_CASE,		references=__SCREAMING_SNAKE_CASE,		)

        snake_case_ :    List[str]     			=						metric.compute()
        # Use accelerator.print to print only on the main process.
        accelerator.print(f'epoch {epoch}:',		__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	):
    """simple docstring"""






    snake_case_ :    Tuple     			=						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.",		)
    # New Code #
    parser.add_argument(
        "--gradient_accumulation_steps",		type=__SCREAMING_SNAKE_CASE,		default=1,		help="The number of minibatches to be ran before gradients are accumulated.",		)
    parser.add_argument("--cpu",		action="store_true",		help="If passed, will train on the CPU."       )
    snake_case_ :    List[str]     			=						parser.parse_args()
    snake_case_ :    Union[str, Any]     			=						{"lr": 2E-5, "num_epochs": 3, "seed": 4_2, "batch_size": 1_6}
    training_function(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )


if __name__ == "__main__":
  main()






'''simple docstring'''
import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
  from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor






class    UpperCAmelCase_				:





 def __init__(    self ,       lowercase_ ,       lowercase_=13 ,       lowercase_=32 ,       lowercase_=3 ,       lowercase_=4 ,       lowercase_=[10, 20, 30, 40] ,       lowercase_=[2, 2, 3, 2] ,       lowercase_=True ,       lowercase_=True ,       lowercase_=37 ,       lowercase_="gelu" ,       lowercase_=10 ,       lowercase_=0.02 ,       lowercase_=["stage2", "stage3", "stage4"] ,       lowercase_=[2, 3, 4] ,       lowercase_=None ,       ):
     snake_case_ :    Dict     			=						parent
     snake_case_ :    int     			=						batch_size
     snake_case_ :    List[str]     			=						image_size
     snake_case_ :    Tuple     			=						num_channels
     snake_case_ :    Dict     			=						num_stages
     snake_case_ :    Optional[int]     			=						hidden_sizes
     snake_case_ :    Optional[Any]     			=						depths
     snake_case_ :    Optional[int]     			=						is_training
     snake_case_ :    Any     			=						use_labels
     snake_case_ :    List[str]     			=						intermediate_size
     snake_case_ :    List[str]     			=						hidden_act
     snake_case_ :    Dict     			=						num_labels
     snake_case_ :    Union[str, Any]     			=						initializer_range
     snake_case_ :    int     			=						out_features
     snake_case_ :    List[str]     			=						out_indices
     snake_case_ :    int     			=						scope





 def        snake_case__		(    self):
     snake_case_ :    Any     			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

     snake_case_ :    Tuple     			=						None
     if self.use_labels:
         snake_case_ :    Any     			=						ids_tensor([self.batch_size] ,       self.num_labels)

     snake_case_ :    int     			=						self.get_config()

     return config, pixel_values, labels





 def        snake_case__		(    self):
     return ConvNextVaConfig(
         num_channels=self.num_channels ,       hidden_sizes=self.hidden_sizes ,       depths=self.depths ,       num_stages=self.num_stages ,       hidden_act=self.hidden_act ,       is_decoder=lowercase_ ,       initializer_range=self.initializer_range ,       out_features=self.out_features ,       out_indices=self.out_indices ,       num_labels=self.num_labels ,       )





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_ ,       lowercase_):
     snake_case_ :    Dict     			=						ConvNextVaModel(config=lowercase_)
     model.to(lowercase_)
     model.eval()
     snake_case_ :    Dict     			=						model(lowercase_)
     # expected last hidden states: B, C, H // 32, W // 32
     self.parent.assertEqual(
         result.last_hidden_state.shape ,       (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32) ,       )





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_ ,       lowercase_):
     snake_case_ :    Tuple     			=						ConvNextVaForImageClassification(lowercase_)
     model.to(lowercase_)
     model.eval()
     snake_case_ :    int     			=						model(lowercase_ ,       labels=lowercase_)
     self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_labels))





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_ ,       lowercase_):
     snake_case_ :    Tuple     			=						ConvNextVaBackbone(config=lowercase_)
     model.to(lowercase_)
     model.eval()
     snake_case_ :    Any     			=						model(lowercase_)

     # verify hidden states
     self.parent.assertEqual(len(result.feature_maps) ,       len(config.out_features))
     self.parent.assertListEqual(list(result.feature_maps[0].shape) ,       [self.batch_size, self.hidden_sizes[1], 4, 4])

     # verify channels
     self.parent.assertEqual(len(model.channels) ,       len(config.out_features))
     self.parent.assertListEqual(model.channels ,       config.hidden_sizes[1:])

     # verify backbone works with out_features=None
     snake_case_ :    Dict     			=						None
     snake_case_ :    str     			=						ConvNextVaBackbone(config=lowercase_)
     model.to(lowercase_)
     model.eval()
     snake_case_ :    List[str]     			=						model(lowercase_)

     # verify feature maps
     self.parent.assertEqual(len(result.feature_maps) ,       1)
     self.parent.assertListEqual(list(result.feature_maps[0].shape) ,       [self.batch_size, self.hidden_sizes[-1], 1, 1])

     # verify channels
     self.parent.assertEqual(len(model.channels) ,       1)
     self.parent.assertListEqual(model.channels ,       [config.hidden_sizes[-1]])





 def        snake_case__		(    self):
     snake_case_ :    Optional[int]     			=						self.prepare_config_and_inputs()
     snake_case_			, snake_case_			, snake_case_ :    str     			=						config_and_inputs
     snake_case_ :    int     			=						{"pixel_values": pixel_values}
     return config, inputs_dict





 def        snake_case__		(    self):
     snake_case_ :    Union[str, Any]     			=						self.prepare_config_and_inputs()
     snake_case_			, snake_case_			, snake_case_ :    Dict     			=						config_and_inputs
     snake_case_ :    Union[str, Any]     			=						{"pixel_values": pixel_values, "labels": labels}
     return config, inputs_dict






@require_torch
class    UpperCAmelCase_				(	snake_case__    ,						snake_case__    ,						unittest.TestCase					):
 UpperCAmelCase_		    =						(
     (
         ConvNextVaModel,
         ConvNextVaForImageClassification,
         ConvNextVaBackbone,
     )
     if is_torch_available()
     else ()
 )
 UpperCAmelCase_		    =						(
     {"""feature-extraction""": ConvNextVaModel, """image-classification""": ConvNextVaForImageClassification}
     if is_torch_available()
     else {}
 )

 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False
 UpperCAmelCase_		    =						False





 def        snake_case__		(    self):
     snake_case_ :    List[Any]     			=						ConvNextVaModelTester(self)
     snake_case_ :    str     			=						ConfigTester(self ,       config_class=lowercase_ ,       has_text_modality=lowercase_ ,       hidden_size=37)





 def        snake_case__		(    self):
     self.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.create_and_test_config_with_num_labels()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()





 def        snake_case__		(    self):
     return





 @unittest.skip(reason="ConvNextV2 does not use inputs_embeds")
 def        snake_case__		(    self):
     pass





 @unittest.skip(reason="ConvNextV2 does not support input and output embeddings")
 def        snake_case__		(    self):
     pass





 @unittest.skip(reason="ConvNextV2 does not use feedforward chunking")
 def        snake_case__		(    self):
     pass





 def        snake_case__		(    self):
     if not self.model_tester.is_training:
         return

     for model_class in self.all_model_classes:
         snake_case_			, snake_case_ :    Optional[Any]     			=						self.model_tester.prepare_config_and_inputs_with_labels()
         snake_case_ :    Union[str, Any]     			=						True

         if model_class.__name__ in [
             *get_values(lowercase_),
             *get_values(lowercase_),
         ]:
             continue

         snake_case_ :    List[Any]     			=						model_class(lowercase_)
         model.to(lowercase_)
         model.train()
         snake_case_ :    int     			=						self._prepare_for_class(lowercase_ ,       lowercase_ ,       return_labels=lowercase_)
         snake_case_ :    Union[str, Any]     			=						model(**lowercase_).loss
         loss.backward()





 def        snake_case__		(    self):
     if not self.model_tester.is_training:
         return

     for model_class in self.all_model_classes:
         snake_case_			, snake_case_ :    List[str]     			=						self.model_tester.prepare_config_and_inputs_with_labels()
         snake_case_ :    Dict     			=						False
         snake_case_ :    Any     			=						True

         if (
             model_class.__name__
             in [*get_values(lowercase_), *get_values(lowercase_)]
             or not model_class.supports_gradient_checkpointing
         ):
             continue

         snake_case_ :    Any     			=						model_class(lowercase_)
         model.to(lowercase_)
         model.gradient_checkpointing_enable()
         model.train()
         snake_case_ :    Any     			=						self._prepare_for_class(lowercase_ ,       lowercase_ ,       return_labels=lowercase_)
         snake_case_ :    Optional[int]     			=						model(**lowercase_).loss
         loss.backward()





 def        snake_case__		(    self):
     snake_case_			, snake_case_ :    Dict     			=						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         snake_case_ :    Any     			=						model_class(lowercase_)
         snake_case_ :    Any     			=						inspect.signature(model.forward)
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         snake_case_ :    List[Any]     			=						[*signature.parameters.keys()]

         snake_case_ :    str     			=						["pixel_values"]
         self.assertListEqual(arg_names[:1] ,       lowercase_)





 def        snake_case__		(    self):
     snake_case_ :    Tuple     			=						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_)





 def        snake_case__		(    self):
     def check_hidden_states_output(lowercase_ ,       lowercase_ ,       lowercase_):
         snake_case_ :    int     			=						model_class(lowercase_)
         model.to(lowercase_)
         model.eval()

         with torch.no_grad():
             snake_case_ :    int     			=						model(**self._prepare_for_class(lowercase_ ,       lowercase_))

         snake_case_ :    int     			=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

         snake_case_ :    Dict     			=						self.model_tester.num_stages
         self.assertEqual(len(lowercase_) ,       expected_num_stages + 1)

         # ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
         self.assertListEqual(
             list(hidden_states[0].shape[-2:]) ,       [self.model_tester.image_size // 4, self.model_tester.image_size // 4] ,       )

     snake_case_			, snake_case_ :    Any     			=						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         snake_case_ :    Tuple     			=						True
         check_hidden_states_output(lowercase_ ,       lowercase_ ,       lowercase_)

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         snake_case_ :    Union[str, Any]     			=						True

         check_hidden_states_output(lowercase_ ,       lowercase_ ,       lowercase_)





 def        snake_case__		(    self):
     snake_case_ :    Union[str, Any]     			=						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_)





 @slow
 def        snake_case__		(    self):
     for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case_ :    Any     			=						ConvNextVaModel.from_pretrained(lowercase_)
         self.assertIsNotNone(lowercase_)


def 						UpperCamelCase_     (	):
    """simple docstring"""






    snake_case_ :    List[Any]     			=						Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
    return image






@require_torch
@require_vision
class    UpperCAmelCase_				(	unittest.TestCase					):





 @cached_property
 def        snake_case__		(    self):
     return AutoImageProcessor.from_pretrained("facebook/convnextv2-tiny-1k-224") if is_vision_available() else None





 @slow
 def        snake_case__		(    self):
     snake_case_ :    List[str]     			=						ConvNextVaForImageClassification.from_pretrained("facebook/convnextv2-tiny-1k-224").to(lowercase_)

     snake_case_ :    List[Any]     			=						self.default_image_processor
     snake_case_ :    str     			=						prepare_img()
     snake_case_ :    Any     			=						preprocessor(images=lowercase_ ,       return_tensors="pt").to(lowercase_)

     # forward pass
     with torch.no_grad():
         snake_case_ :    List[str]     			=						model(**lowercase_)

     # verify the logits
     snake_case_ :    str     			=						torch.Size((1, 10_00))
     self.assertEqual(outputs.logits.shape ,       lowercase_)

     snake_case_ :    str     			=						torch.tensor([0.9_996, 0.1_966, -0.4_386]).to(lowercase_)
     self.assertTrue(torch.allclose(outputs.logits[0, :3] ,       lowercase_ ,       atol=1E-4))




