




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


A     =16
A     =32





def 						snake_case_							(_a	:							Accelerator    ,   _a	:							DatasetDict    ,   _a	:							List[int]    ,   _a	:							List[int]    ,   _a	:							int = 1_6     ):
    UpperCAmelCase			       =		AutoTokenizer.from_pretrained('''bert-base-cased'''     )
    UpperCAmelCase			       =		DatasetDict(
        {
            '''train''': dataset['''train'''].select(_a     ),
            '''validation''': dataset['''train'''].select(_a     ),
            '''test''': dataset['''validation'''],
        }     )

    def tokenize_function(_a	:							int     ):
        # max_length=None => use the model max length (it's actually the default)
        UpperCAmelCase			       =		tokenizer(examples['''sentence1''']    ,   examples['''sentence2''']    ,   truncation=_a    ,   max_length=_a     )
        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(
            _a    ,   batched=_a    ,   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(_a	:							List[str]     ):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        UpperCAmelCase			       =		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":
            UpperCAmelCase			       =		1_6
        elif accelerator.mixed_precision != "no":
            UpperCAmelCase			       =		8
        else:
            UpperCAmelCase			       =		None

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

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

    UpperCAmelCase			       =		DataLoader(
        tokenized_datasets['''test''']    ,   shuffle=_a    ,   collate_fn=_a    ,   batch_size=_a     )

    return train_dataloader, eval_dataloader, test_dataloader





def 						snake_case_							(_a	:							Any    ,   _a	:							Union[str, Any]     ):
    # New Code #
    UpperCAmelCase			       =		[]
    # Download the dataset
    UpperCAmelCase			       =		load_dataset('''glue'''    ,   '''mrpc'''     )
    # Create our splits
    UpperCAmelCase			       =		StratifiedKFold(n_splits=int(args.num_folds     )     )
    # 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(_a     )

    # New Code #
    # Create our folds:
    UpperCAmelCase			       =		kfold.split(np.zeros(datasets['''train'''].num_rows     )    ,   datasets['''train''']['''label''']     )
    UpperCAmelCase			       =		[]
    # Iterate over them
    for i, (train_idxs, valid_idxs) in enumerate(_a     ):
        UpperCAmelCase       ,     UpperCAmelCase       ,     UpperCAmelCase			       =		get_fold_dataloaders(
            _a    ,   _a    ,   _a    ,   _a    ,   )
        # 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=_a     )

        # 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=_a     )

        # Instantiate scheduler
        UpperCAmelCase			       =		get_linear_schedule_with_warmup(
            optimizer=_a    ,   num_warmup_steps=1_0_0    ,   num_training_steps=(len(_a     ) * 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(
            _a    ,   _a    ,   _a    ,   _a    ,   _a     )

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

            model.eval()
            for step, batch in enumerate(_a     ):
                # We could avoid this line since we set the accelerator with `device_placement=True`.
                batch.to(accelerator.device     )
                with torch.no_grad():
                    UpperCAmelCase			       =		model(**_a     )
                UpperCAmelCase			       =		outputs.logits.argmax(dim=-1     )
                UpperCAmelCase       ,     UpperCAmelCase			       =		accelerator.gather_for_metrics((predictions, batch['''labels'''])     )
                metric.add_batch(
                    predictions=_a    ,   references=_a    ,   )

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

        # New Code #
        # We also run predictions on the test set at the very end
        UpperCAmelCase			       =		[]
        for step, batch in enumerate(_a     ):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            batch.to(accelerator.device     )
            with torch.no_grad():
                UpperCAmelCase			       =		model(**_a     )
            UpperCAmelCase			       =		outputs.logits
            UpperCAmelCase       ,     UpperCAmelCase			       =		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(_a    ,   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:
    UpperCAmelCase			       =		torch.cat(_a    ,   dim=0     )
    UpperCAmelCase			       =		torch.stack(_a    ,   dim=0     ).sum(dim=0     ).div(int(args.num_folds     )     ).argmax(dim=-1     )
    UpperCAmelCase			       =		metric.compute(predictions=_a    ,   references=_a     )
    accelerator.print('''Average test metrics from all folds:'''    ,   _a     )





def 						snake_case_							():
    UpperCAmelCase			       =		argparse.ArgumentParser(description='''Simple example of training script.'''     )
    parser.add_argument(
        '''--mixed_precision'''    ,   type=_a    ,   default=_a    ,   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=_a    ,   default=3    ,   help='''The number of splits to perform across the dataset'''     )
    UpperCAmelCase			       =		parser.parse_args()
    UpperCAmelCase			       =		{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 4_2, '''batch_size''': 1_6}
    training_function(_a    ,   _a     )


if __name__ == "__main__":
   main()







import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict

@pytest.mark.parametrize(
    """files"""    ,      [
        ["""full:README.md""", """dataset_infos.json"""],
        ["""empty:README.md""", """dataset_infos.json"""],
        ["""dataset_infos.json"""],
        ["""full:README.md"""],
    ]    ,      )
def 						UpperCamelCase					(lowercase_: List[str]    ,      lowercase_: Optional[int]					)     ->      List[str]:
	A__							:    Tuple        =		tmp_path_factory.mktemp("""dset_infos_dir"""					)
	if "full:README.md" in files:
		with open(dataset_infos_dir / """README.md"""    ,      """w"""					) as f:
			f.write("""---\ndataset_info:\n  dataset_size: 42\n---"""					)
	if "empty:README.md" in files:
		with open(dataset_infos_dir / """README.md"""    ,      """w"""					) as f:
			f.write(""""""					)
    # we want to support dataset_infos.json for backward compatibility
	if "dataset_infos.json" in files:
		with open(dataset_infos_dir / """dataset_infos.json"""    ,      """w"""					) as f:
			f.write("""{\"default\": {\"dataset_size\": 42}}"""					)
	A__							:    List[Any]        =		DatasetInfosDict.from_directory(lowercase_					)
	assert dataset_infos
	assert dataset_infos["default"].dataset_size == 42

@pytest.mark.parametrize(
    """dataset_info"""    ,      [
        DatasetInfo(),
        DatasetInfo(
            description="""foo"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train"""}]    ,      download_size=42    ,      ),
    ]    ,      )
def 						UpperCamelCase					(lowercase_: str    ,      lowercase_: DatasetInfo					)     ->      List[Any]:
	A__							:    Union[str, Any]        =		str(lowercase_					)
	dataset_info.write_to_directory(lowercase_					)
	A__							:    List[Any]        =		DatasetInfo.from_directory(lowercase_					)
	assert dataset_info == reloaded
	assert os.path.exists(os.path.join(lowercase_    ,      """dataset_info.json"""					)					)

def 						UpperCamelCase					()     ->      List[Any]:
	A__							:    Union[str, Any]        =		DatasetInfo(
	    description="""foo"""    ,      citation="""bar"""    ,      homepage="""https://foo.bar"""    ,      license="""CC0"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      post_processed={}    ,      supervised_keys=()    ,      task_templates=[]    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train""", """num_examples""": 42}]    ,      download_checksums={}    ,      download_size=1337    ,      post_processing_size=442    ,      dataset_size=1234    ,      size_in_bytes=1337 + 442 + 1234    ,      )
	A__							:    Dict        =		dataset_info._to_yaml_dict()
	assert sorted(lowercase_					) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML					)
	for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
		assert key in dataset_info_yaml_dict
		assert isinstance(dataset_info_yaml_dict[key]    ,      (list, dict, int, str)					)
	A__							:    Union[str, Any]        =		yaml.safe_dump(lowercase_					)
	A__							:    List[Any]        =		yaml.safe_load(lowercase_					)
	assert dataset_info_yaml_dict == reloaded

def 						UpperCamelCase					()     ->      List[str]:
	A__							:    Optional[int]        =		DatasetInfo()
	A__							:    List[Any]        =		dataset_info._to_yaml_dict()
	assert dataset_info_yaml_dict == {}


@pytest.mark.parametrize(
    """dataset_infos_dict"""    ,      [
        DatasetInfosDict(),
        DatasetInfosDict({"""default""": DatasetInfo()}					),
        DatasetInfosDict({"""my_config_name""": DatasetInfo()}					),
        DatasetInfosDict(
            {
                """default""": DatasetInfo(
                    description="""foo"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train"""}]    ,      download_size=42    ,      )
            }					),
        DatasetInfosDict(
            {
                """v1""": DatasetInfo(dataset_size=42					),
                """v2""": DatasetInfo(dataset_size=1337					),
            }					),
    ]    ,      )
def 						UpperCamelCase					(lowercase_: Tuple    ,      lowercase_: DatasetInfosDict					)     ->      Optional[Any]:
	A__							:    List[Any]        =		str(lowercase_					)
	dataset_infos_dict.write_to_directory(lowercase_					)
	A__							:    Dict        =		DatasetInfosDict.from_directory(lowercase_					)

	# the config_name of the dataset_infos_dict take over the attribute
	for config_name, dataset_info in dataset_infos_dict.items():
		A__							:    Optional[int]        =		config_name
		# the yaml representation doesn't include fields like description or citation
		# so we just test that we can recover what we can from the yaml
		A__							:    List[str]        =		DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()					)
	assert dataset_infos_dict == reloaded

	if dataset_infos_dict:
		assert os.path.exists(os.path.join(lowercase_    ,      """README.md"""					)					)



