




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


A     =get_tests_dir('fixtures/dummy-config.json')


class 			_a  (	unittest.TestCase ):
       def     A      (	self						:	Tuple				):



           '''simple docstring'''

           UpperCAmelCase			       =		0
       def     A      (	self						:	Any				):



           '''simple docstring'''

           self.assertIsNotNone(transformers.models.auto.__spec__				)
           self.assertIsNotNone(importlib.util.find_spec('''transformers.models.auto'''				)				)
       def     A      (	self						:	int				):



           '''simple docstring'''

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



           '''simple docstring'''

           UpperCAmelCase			       =		AutoConfig.from_pretrained(lowercase_				)
           self.assertIsInstance(lowercase_   ,				lowercase_				)
       def     A      (	self						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		AutoConfig.from_pretrained(lowercase_				)
           self.assertIsInstance(lowercase_   ,				lowercase_				)
       def     A      (	self						:	Tuple				):



           '''simple docstring'''

           UpperCAmelCase			       =		AutoConfig.for_model('''roberta'''				)
           self.assertIsInstance(lowercase_   ,				lowercase_				)
       def     A      (	self						:	Union[str, Any]				):



           '''simple docstring'''

           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     A      (	self						:	Tuple				):



           '''simple docstring'''

           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     A      (	self						:	int				):



           '''simple docstring'''

           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     A      (	self						:	int				):



           '''simple docstring'''

           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     A      (	self						:	Tuple				):



           '''simple docstring'''

           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     A      (	self						:	int				):



           '''simple docstring'''

           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     A      (	self						:	List[Any]				):



           '''simple docstring'''

           class 			_a  (	UpperCamelCase__ ):
                  __a				:     List[str]   				=					"""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"]











'''simple docstring'''

import argparse
from typing import List

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset

# New Code #
# We'll be using StratifiedKFold for this example
from sklearn.model_selection import StratifiedKFold
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 perform Cross Validation,
# 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
#
########################################################################


__lowerCAmelCase							:	Optional[int]			=16
__lowerCAmelCase							:	Tuple			=32





def 						UpperCamelCase				(							_lowerCamelCase					:					Accelerator						,       _lowerCamelCase					:					DatasetDict						,       _lowerCamelCase					:					List[int]						,       _lowerCamelCase					:					List[int]						,       _lowerCamelCase					:					int = 16       ):
    A__				=			AutoTokenizer.from_pretrained("bert-base-cased"       )
    A__				=			DatasetDict(
        {
            "train": dataset["train"].select(_lowerCamelCase       ),
            "validation": dataset["train"].select(_lowerCamelCase       ),
            "test": dataset["validation"],
        }       )

    def tokenize_function(_lowerCamelCase					:					Dict       ):
        # max_length=None => use the model max length (it's actually the default)
        A__				=			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
    # starting with the main process first:
    with accelerator.main_process_first():
        A__				=			datasets.map(
            _lowerCamelCase						,       batched=_lowerCamelCase						,       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
    A__				=			tokenized_datasets.rename_column("label"						,       "labels"       )

    def collate_fn(_lowerCamelCase					:					Any       ):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        A__				=			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":
            A__				=			16
        elif accelerator.mixed_precision != "no":
            A__				=			8
        else:
            A__				=			None

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

    # Instantiate dataloaders.
    A__				=			DataLoader(
        tokenized_datasets["train"]						,       shuffle=_lowerCamelCase						,       collate_fn=_lowerCamelCase						,       batch_size=_lowerCamelCase       )
    A__				=			DataLoader(
        tokenized_datasets["validation"]						,       shuffle=_lowerCamelCase						,       collate_fn=_lowerCamelCase						,       batch_size=_lowerCamelCase       )

    A__				=			DataLoader(
        tokenized_datasets["test"]						,       shuffle=_lowerCamelCase						,       collate_fn=_lowerCamelCase						,       batch_size=_lowerCamelCase       )

    return train_dataloader, eval_dataloader, test_dataloader





def 						UpperCamelCase				(							_lowerCamelCase					:					List[Any]						,       _lowerCamelCase					:					str       ):
    # New Code #
    A__				=			[]
    # Download the dataset
    A__				=			load_dataset("glue"						,       "mrpc"       )
    # Create our splits
    A__				=			StratifiedKFold(n_splits=int(args.num_folds       )       )
    # Initialize accelerator
    A__				=			Accelerator(cpu=args.cpu						,       mixed_precision=args.mixed_precision       )
    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    A__				=			config["lr"]
    A__				=			int(config["num_epochs"]       )
    A__				=			int(config["seed"]       )
    A__				=			int(config["batch_size"]       )

    A__				=			evaluate.load("glue"						,       "mrpc"       )

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

    set_seed(_lowerCamelCase       )

    # New Code #
    # Create our folds:
    A__				=			kfold.split(np.zeros(datasets["train"].num_rows       )						,       datasets["train"]["label"]       )
    A__				=			[]
    # Iterate over them
    for i, (train_idxs, valid_idxs) in enumerate(_lowerCamelCase       ):
        A__,	A__,	A__				=			get_fold_dataloaders(
            _lowerCamelCase						,       _lowerCamelCase						,       _lowerCamelCase						,       _lowerCamelCase						,       )
        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
        A__				=			AutoModelForSequenceClassification.from_pretrained("bert-base-cased"						,       return_dict=_lowerCamelCase       )

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

        # Instantiate optimizer
        A__				=			AdamW(params=model.parameters()						,       lr=_lowerCamelCase       )

        # Instantiate scheduler
        A__				=			get_linear_schedule_with_warmup(
            optimizer=_lowerCamelCase						,       num_warmup_steps=1_00						,       num_training_steps=(len(_lowerCamelCase       ) * 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.
        A__,	A__,	A__,	A__,	A__				=			accelerator.prepare(
            _lowerCamelCase						,       _lowerCamelCase						,       _lowerCamelCase						,       _lowerCamelCase						,       _lowerCamelCase       )

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

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

            A__				=			metric.compute()
            # Use accelerator.print to print only on the main process.
            accelerator.print(F"epoch {epoch}:"						,       _lowerCamelCase       )

        # New Code #
        # We also run predictions on the test set at the very end
        A__				=			[]
        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():
                A__				=			model(**_lowerCamelCase       )
            A__				=			outputs.logits
            A__,	A__				=			accelerator.gather_for_metrics((predictions, batch["labels"])       )
            fold_predictions.append(predictions.cpu()       )
            if i == 0:
                # We need all of the test predictions
                test_references.append(references.cpu()       )
        # Use accelerator.print to print only on the main process.
        test_predictions.append(torch.cat(_lowerCamelCase						,       dim=0       )       )
        # We now need to release all our memory and get rid of the current model, optimizer, etc
        accelerator.free_memory()
    # New Code #
    # Finally we check the accuracy of our folded results:
    A__				=			torch.cat(_lowerCamelCase						,       dim=0       )
    A__				=			torch.stack(_lowerCamelCase						,       dim=0       ).sum(dim=0       ).div(int(args.num_folds       )       ).argmax(dim=-1       )
    A__				=			metric.compute(predictions=_lowerCamelCase						,       references=_lowerCamelCase       )
    accelerator.print("Average test metrics from all folds:"						,       _lowerCamelCase       )





def 						UpperCamelCase				(							):
    A__				=			argparse.ArgumentParser(description="Simple example of training script."       )
    parser.add_argument(
        "--mixed_precision"						,       type=_lowerCamelCase						,       default=_lowerCamelCase						,       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."       )
    # New Code #
    parser.add_argument("--num_folds"						,       type=_lowerCamelCase						,       default=3						,       help="The number of splits to perform across the dataset"       )
    A__				=			parser.parse_args()
    A__				=			{"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
    training_function(_lowerCamelCase						,       _lowerCamelCase       )


if __name__ == "__main__":
  main()



