





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


lowerCAmelCase							:       List[Any]					 =	16
lowerCAmelCase							:       Tuple					 =	32






def 		lowerCAmelCase  (	UpperCamelCase__   :   Accelerator       ,			UpperCamelCase__   :   int = 16					)      ->    str:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   AutoTokenizer.from_pretrained('''bert-base-cased'''					)
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   load_dataset('''glue'''       ,			'''mrpc'''					)

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

						def collate_fn(UpperCamelCase__   :   Any					):
												# On TPU it's best to pad everything to the same length or training will be very slow.
												__SCREAMING_SNAKE_CASE:				int           =   128 if accelerator.distributed_type == DistributedType.TPU else None
												# When using mixed precision we want round multiples of 8/16
												if accelerator.mixed_precision == "fp8":
																		__SCREAMING_SNAKE_CASE:				Tuple           =   16
												elif accelerator.mixed_precision != "no":
																		__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   8
												else:
																		__SCREAMING_SNAKE_CASE:				int           =   None

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

						# Instantiate dataloaders.
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   DataLoader(
						    tokenized_datasets['''train''']       ,			shuffle=UpperCamelCase__       ,			collate_fn=UpperCamelCase__       ,			batch_size=UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				List[Any]           =   DataLoader(
						    tokenized_datasets['''validation''']       ,			shuffle=UpperCamelCase__       ,			collate_fn=UpperCamelCase__       ,			batch_size=UpperCamelCase__					)

						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

							lowerCAmelCase							:       Optional[int]					 =	mocked_dataloaders  # noqa: F811






def 		lowerCAmelCase  (	UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   Tuple					)      ->    Optional[Any]:

						"""simple docstring"""





						if os.environ.get('''TESTING_MOCKED_DATALOADERS'''       ,			UpperCamelCase__					) == "1":
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   2
						# New Code #
						__SCREAMING_SNAKE_CASE:				int           =   int(args.gradient_accumulation_steps					)
						# Initialize accelerator
						__SCREAMING_SNAKE_CASE:				List[Any]           =   Accelerator(
						    cpu=args.cpu       ,			mixed_precision=args.mixed_precision       ,			gradient_accumulation_steps=UpperCamelCase__					)
						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
						__SCREAMING_SNAKE_CASE:				List[Any]           =   config['''lr''']
						__SCREAMING_SNAKE_CASE:				List[Any]           =   int(config['''num_epochs''']					)
						__SCREAMING_SNAKE_CASE:				Any           =   int(config['''seed''']					)
						__SCREAMING_SNAKE_CASE:				List[str]           =   int(config['''batch_size''']					)

						__SCREAMING_SNAKE_CASE:				Tuple           =   evaluate.load('''glue'''       ,			'''mrpc'''					)

						set_seed(UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				str           =   get_dataloaders(UpperCamelCase__       ,			UpperCamelCase__					)
						# Instantiate the model (we build the model here so that the seed also control new weights initialization)
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''       ,			return_dict=UpperCamelCase__					)

						# 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).
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   model.to(accelerator.device					)

						# Instantiate optimizer
						__SCREAMING_SNAKE_CASE:				str           =   AdamW(params=model.parameters()       ,			lr=UpperCamelCase__					)

						# Instantiate scheduler
						__SCREAMING_SNAKE_CASE:				int           =   get_linear_schedule_with_warmup(
						    optimizer=UpperCamelCase__       ,			num_warmup_steps=100       ,			num_training_steps=(len(UpperCamelCase__					) * 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.
						__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				Any           =   accelerator.prepare(
						    UpperCamelCase__       ,			UpperCamelCase__       ,			UpperCamelCase__       ,			UpperCamelCase__       ,			UpperCamelCase__					)

						# Now we train the model
						for epoch in range(UpperCamelCase__					):
												model.train()
												for step, batch in enumerate(UpperCamelCase__					):
																		# 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(UpperCamelCase__					):
																								__SCREAMING_SNAKE_CASE:				List[Any]           =   model(**UpperCamelCase__					)
																								__SCREAMING_SNAKE_CASE:				Any           =   output.loss
																								accelerator.backward(UpperCamelCase__					)
																								optimizer.step()
																								lr_scheduler.step()
																								optimizer.zero_grad()

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

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






def 		lowerCAmelCase  (	)      ->    Tuple:

						"""simple docstring"""





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


if __name__ == "__main__":
							main()










import os
from pathlib import Path
from unittest.mock import patch

import pytest
import zstandard as zstd

from datasets.download.download_config import DownloadConfig
from datasets.utils.file_utils import (
    OfflineModeIsEnabled,
    cached_path,
    fsspec_get,
    fsspec_head,
    ftp_get,
    ftp_head,
    get_from_cache,
    http_get,
    http_head,
)


lowerCAmelCase							:       Dict					 =	"""\
    Text data.
    Second line of data."""

lowerCAmelCase							:       str					 =	"""file"""






@pytest.fixture(scope='''session'''					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Union[str, Any]					)      ->    Any:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Any           =   tmp_path_factory.mktemp('''data'''					) / (FILE_PATH + '''.zstd''')
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   bytes(UpperCamelCase__       ,			'''utf-8'''					)
						with zstd.open(UpperCamelCase__       ,			'''wb'''					) as f:
												f.write(UpperCamelCase__					)
						return path






@pytest.fixture
def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[Any]					)      ->    int:

						"""simple docstring"""





						with open(os.path.join(tmpfs.local_root_dir       ,			UpperCamelCase__					)       ,			'''w'''					) as f:
												f.write(UpperCamelCase__					)
						return FILE_PATH






@pytest.mark.parametrize('''compression_format'''       ,			['''gzip''', '''xz''', '''zstd''']					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   Dict       ,			UpperCamelCase__   :   Optional[Any]       ,			UpperCamelCase__   :   Dict       ,			UpperCamelCase__   :   List[Any]					)      ->    Any:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Dict           =   {'''gzip''': gz_file, '''xz''': xz_file, '''zstd''': zstd_path}
						__SCREAMING_SNAKE_CASE:				List[str]           =   input_paths[compression_format]
						__SCREAMING_SNAKE_CASE:				List[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				List[str]           =   DownloadConfig(cache_dir=UpperCamelCase__       ,			extract_compressed_file=UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				Tuple           =   cached_path(UpperCamelCase__       ,			download_config=UpperCamelCase__					)
						with open(UpperCamelCase__					) as f:
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   f.read()
						with open(UpperCamelCase__					) as f:
												__SCREAMING_SNAKE_CASE:				Dict           =   f.read()
						assert extracted_file_content == expected_file_content






@pytest.mark.parametrize('''default_extracted'''       ,			[True, False]					)
@pytest.mark.parametrize('''default_cache_dir'''       ,			[True, False]					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Tuple       ,			UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   str       ,			UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   Optional[Any]					)      ->    Optional[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Tuple           =   '''custom_cache'''
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   '''custom_extracted_dir'''
						__SCREAMING_SNAKE_CASE:				List[Any]           =   tmp_path / '''custom_extracted_path'''
						if default_extracted:
												__SCREAMING_SNAKE_CASE:				int           =   ('''downloads''' if default_cache_dir else custom_cache_dir, '''extracted''')
						else:
												monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_DIR'''       ,			UpperCamelCase__					)
												monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_PATH'''       ,			str(UpperCamelCase__					)					)
												__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   custom_extracted_path.parts[-2:] if default_cache_dir else (custom_cache_dir, custom_extracted_dir)

						__SCREAMING_SNAKE_CASE:				Dict           =   xz_file
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   (
						    DownloadConfig(extract_compressed_file=UpperCamelCase__					)
						    if default_cache_dir
						    else DownloadConfig(cache_dir=tmp_path / custom_cache_dir       ,			extract_compressed_file=UpperCamelCase__					)
						)
						__SCREAMING_SNAKE_CASE:				int           =   cached_path(UpperCamelCase__       ,			download_config=UpperCamelCase__					)
						assert Path(UpperCamelCase__					).parent.parts[-2:] == expected






def 		lowerCAmelCase  (	UpperCamelCase__   :   Dict					)      ->    Union[str, Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   str(Path(UpperCamelCase__					).resolve()					)
						assert cached_path(UpperCamelCase__					) == text_file
						# relative path
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   str(Path(UpperCamelCase__					).resolve().relative_to(Path(os.getcwd()					)					)					)
						assert cached_path(UpperCamelCase__					) == text_file






def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[Any]					)      ->    Tuple:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[Any]           =   str(tmp_path.resolve() / '''__missing_file__.txt'''					)
						with pytest.raises(UpperCamelCase__					):
												cached_path(UpperCamelCase__					)
						# relative path
						__SCREAMING_SNAKE_CASE:				int           =   '''./__missing_file__.txt'''
						with pytest.raises(UpperCamelCase__					):
												cached_path(UpperCamelCase__					)






def 		lowerCAmelCase  (	UpperCamelCase__   :   Tuple					)      ->    Union[str, Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Tuple           =   get_from_cache(F"""tmp://{tmpfs_file}"""					)
						with open(UpperCamelCase__					) as f:
												__SCREAMING_SNAKE_CASE:				Tuple           =   f.read()
						assert output_file_content == FILE_CONTENT






@patch('''datasets.config.HF_DATASETS_OFFLINE'''       ,			UpperCamelCase__					)
def 		lowerCAmelCase  (	)      ->    Any:

						"""simple docstring"""





						with pytest.raises(UpperCamelCase__					):
												cached_path('''https://huggingface.co'''					)






@patch('''datasets.config.HF_DATASETS_OFFLINE'''       ,			UpperCamelCase__					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   int					)      ->    Union[str, Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[int]           =   tmp_path_factory.mktemp('''data'''					) / '''file.html'''
						with pytest.raises(UpperCamelCase__					):
												http_get('''https://huggingface.co'''       ,			temp_file=UpperCamelCase__					)
						with pytest.raises(UpperCamelCase__					):
												http_head('''https://huggingface.co'''					)






@patch('''datasets.config.HF_DATASETS_OFFLINE'''       ,			UpperCamelCase__					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Dict					)      ->    List[str]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[Any]           =   tmp_path_factory.mktemp('''data'''					) / '''file.html'''
						with pytest.raises(UpperCamelCase__					):
												ftp_get('''ftp://huggingface.co'''       ,			temp_file=UpperCamelCase__					)
						with pytest.raises(UpperCamelCase__					):
												ftp_head('''ftp://huggingface.co'''					)






@patch('''datasets.config.HF_DATASETS_OFFLINE'''       ,			UpperCamelCase__					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[int]					)      ->    Tuple:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Tuple           =   tmp_path_factory.mktemp('''data'''					) / '''file.html'''
						with pytest.raises(UpperCamelCase__					):
												fsspec_get('''s3://huggingface.co'''       ,			temp_file=UpperCamelCase__					)
						with pytest.raises(UpperCamelCase__					):
												fsspec_head('''s3://huggingface.co'''					)



