


import argparse
import json
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
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


SCREAMING_SNAKE_CASE							      =				16
SCREAMING_SNAKE_CASE							      =				32



def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_ = 16							,				lowerCamelCase_ = "bert-base-cased"				)					->		Any:
    _lowercase :							Union[str, Any]   	=	AutoTokenizer.from_pretrained(a_				)
    _lowercase :							Optional[int]   	=	load_dataset('glue'							,				'mrpc'				)

    def tokenize_function(lowerCamelCase_				):
        # max_length=None => use the model max length (it's actually the default)
        _lowercase :							Tuple   	=	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
    _lowercase :							Optional[int]   	=	datasets.map(
        a_							,				batched=a_							,				remove_columns=['idx', 'sentence1', 'sentence2']							,				load_from_cache_file=a_				)

    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
    # transformers library
    _lowercase :							Any   	=	tokenized_datasets.rename_column('label'							,				'labels'				)

    def collate_fn(lowerCamelCase_				):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        if accelerator.distributed_type == DistributedType.TPU:
            return tokenizer.pad(a_							,				padding='max_length'							,				max_length=128							,				return_tensors='pt'				)
        return tokenizer.pad(a_							,				padding='longest'							,				return_tensors='pt'				)

    # Instantiate dataloaders.
    _lowercase :							int   	=	DataLoader(
        tokenized_datasets['train']							,				shuffle=a_							,				collate_fn=a_							,				batch_size=a_				)
    _lowercase :							int   	=	DataLoader(
        tokenized_datasets['validation']							,				shuffle=a_							,				collate_fn=a_							,				batch_size=a_				)

    return train_dataloader, eval_dataloader



def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_				)					->		List[Any]:
    # Initialize accelerator
    _lowercase :							str   	=	Accelerator()

    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    _lowercase :							Optional[int]   	=	config['''lr''']
    _lowercase :							Optional[int]   	=	int(config['num_epochs']				)
    _lowercase :							Dict   	=	int(config['seed']				)
    _lowercase :							str   	=	int(config['batch_size']				)
    _lowercase :							List[str]   	=	args.model_name_or_path

    set_seed(a_				)
    _lowercase :							List[Any]   	=	get_dataloaders(a_							,				a_							,				a_				)

    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    _lowercase :							List[str]   	=	AutoModelForSequenceClassification.from_pretrained(a_							,				return_dict=a_				)

    # Instantiate optimizer
    _lowercase :							Tuple   	=	(
        AdamW
        if accelerator.state.deepspeed_plugin is None
        or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
        else DummyOptim
    )
    _lowercase :							Any   	=	optimizer_cls(params=model.parameters()							,				lr=a_				)

    if accelerator.state.deepspeed_plugin is not None:
        _lowercase :							int   	=	accelerator.state.deepspeed_plugin.deepspeed_config[
            '''gradient_accumulation_steps'''
        ]
    else:
        _lowercase :							Dict   	=	1
    _lowercase :							List[Any]   	=	(len(a_				) * num_epochs) // gradient_accumulation_steps

    # Instantiate scheduler
    if (
        accelerator.state.deepspeed_plugin is None
        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    ):
        _lowercase :							Optional[int]   	=	get_linear_schedule_with_warmup(
            optimizer=a_							,				num_warmup_steps=0							,				num_training_steps=a_							,				)
    else:
        _lowercase :							Optional[int]   	=	DummyScheduler(a_							,				total_num_steps=a_							,				warmup_num_steps=0				)

    # 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.
    _lowercase :							Tuple   	=	accelerator.prepare(
        a_							,				a_							,				a_							,				a_							,				a_				)

    # We need to keep track of how many total steps we have iterated over
    _lowercase :							Tuple   	=	0
    # We also need to keep track of the stating epoch so files are named properly
    _lowercase :							Tuple   	=	0

    # Now we train the model
    _lowercase :							Tuple   	=	evaluate.load('glue'							,				'mrpc'				)
    _lowercase :							Optional[Any]   	=	0
    _lowercase :							Tuple   	=	{}
    for epoch in range(a_							,				a_				):
        model.train()
        for step, batch in enumerate(a_				):
            _lowercase :							Any   	=	model(**a_				)
            _lowercase :							int   	=	outputs.loss
            _lowercase :							List[Any]   	=	loss / gradient_accumulation_steps
            accelerator.backward(a_				)
            if step % gradient_accumulation_steps == 0:
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            overall_step += 1

        model.eval()
        _lowercase :							str   	=	0
        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():
                _lowercase :							int   	=	model(**a_				)
            _lowercase :							List[Any]   	=	outputs.logits.argmax(dim=-1				)
            # It is slightly faster to call this once, than multiple times
            _lowercase :							Optional[Any]   	=	accelerator.gather(
                (predictions, batch['labels'])				)  # If we are in a multiprocess environment, the last batch has duplicates
            if accelerator.use_distributed:
                if step == len(a_				) - 1:
                    _lowercase :							Any   	=	predictions[: len(eval_dataloader.dataset				) - samples_seen]
                    _lowercase :							Any   	=	references[: len(eval_dataloader.dataset				) - samples_seen]
                else:
                    samples_seen += references.shape[0]
            metric.add_batch(
                predictions=a_							,				references=a_							,				)

        _lowercase :							Tuple   	=	metric.compute()
        # Use accelerator.print to print only on the main process.
        accelerator.print(F'''epoch {epoch}:'''							,				a_				)
        _lowercase :							int   	=	eval_metric['''accuracy''']

        if best_performance < eval_metric["accuracy"]:
            _lowercase :							Dict   	=	eval_metric['''accuracy''']

    if args.performance_lower_bound is not None:
        assert (
            args.performance_lower_bound <= best_performance
        ), F'''Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}'''

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        with open(os.path.join(args.output_dir							,				'all_results.json'				)							,				'w'				) as f:
            json.dump(a_							,				a_				)



def        UpperCamelCase_(   )					->		Union[str, Any]:
    _lowercase :							str   	=	argparse.ArgumentParser(description='Simple example of training script tracking peak GPU memory usage.'				)
    parser.add_argument(
        '--model_name_or_path'							,				type=a_							,				default='bert-base-cased'							,				help='Path to pretrained model or model identifier from huggingface.co/models.'							,				required=a_							,				)
    parser.add_argument(
        '--output_dir'							,				type=a_							,				default='.'							,				help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'							,				)
    parser.add_argument(
        '--performance_lower_bound'							,				type=a_							,				default=a_							,				help='Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.'							,				)
    parser.add_argument(
        '--num_epochs'							,				type=a_							,				default=3							,				help='Number of train epochs.'							,				)
    _lowercase :							Optional[Any]   	=	parser.parse_args()
    _lowercase :							Optional[int]   	=	{'''lr''': 2e-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}
    training_function(a_							,				a_				)


if __name__ == "__main__":
 main()






import subprocess
import sys

from transformers import BertConfig, BertModel, BertTokenizer, pipeline
from transformers.testing_utils import TestCasePlus, require_torch






class _lowerCamelCase(		_a     ):
       @require_torch
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

           _lowercase :							Union[str, Any]   	=	'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

           _lowercase :							List[str]   	=	'\nimport socket\ndef offline_socket(*args, **kwargs): raise RuntimeError("Offline mode is enabled, we shouldn\'t access internet")\nsocket.socket = offline_socket\n        '

           # Force fetching the files so that we can use the cache
           _lowercase :							Union[str, Any]   	=	'hf-internal-testing/tiny-random-bert'
           BertConfig.from_pretrained(lowerCamelCase)
           BertModel.from_pretrained(lowerCamelCase)
           BertTokenizer.from_pretrained(lowerCamelCase)
           pipeline(task='fill-mask',   model=lowerCamelCase)

           # baseline - just load from_pretrained with normal network
           _lowercase :							Any   	=	[sys.executable, '-c', '\n'.join([load, run, mock])]

           # should succeed
           _lowercase :							List[str]   	=	self.get_env()
           # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
           _lowercase :							Any   	=	'1'
           _lowercase :							List[Any]   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())
       @require_torch
       def     UpperCamelCase  (      self) ->     List[Any]:




           """simple docstring"""

           _lowercase :							int   	=	'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

           _lowercase :							Union[str, Any]   	=	'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

           _lowercase :							List[Any]   	=	'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Faking flaky internet")\nsocket.socket = offline_socket\n        '

           # Force fetching the files so that we can use the cache
           _lowercase :							Any   	=	'hf-internal-testing/tiny-random-bert'
           BertConfig.from_pretrained(lowerCamelCase)
           BertModel.from_pretrained(lowerCamelCase)
           BertTokenizer.from_pretrained(lowerCamelCase)
           pipeline(task='fill-mask',   model=lowerCamelCase)

           # baseline - just load from_pretrained with normal network
           _lowercase :							List[str]   	=	[sys.executable, '-c', '\n'.join([load, run, mock])]

           # should succeed
           _lowercase :							Any   	=	self.get_env()
           _lowercase :							str   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())
       @require_torch
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase :							Tuple   	=	'\nfrom transformers import BertConfig, BertModel, BertTokenizer\n        '

           _lowercase :							int   	=	'\nmname = "hf-internal-testing/tiny-random-bert-sharded"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nprint("success")\n        '

           _lowercase :							List[Any]   	=	'\nimport socket\ndef offline_socket(*args, **kwargs): raise ValueError("Offline mode is enabled")\nsocket.socket = offline_socket\n        '

           # baseline - just load from_pretrained with normal network
           _lowercase :							int   	=	[sys.executable, '-c', '\n'.join([load, run])]

           # should succeed
           _lowercase :							List[str]   	=	self.get_env()
           _lowercase :							Dict   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())

           # next emulate no network
           _lowercase :							Union[str, Any]   	=	[sys.executable, '-c', '\n'.join([load, mock, run])]

           # Doesn't fail anymore since the model is in the cache due to other tests, so commenting this.
           # env["TRANSFORMERS_OFFLINE"] = "0"
           # result = subprocess.run(cmd, env=env, check=False, capture_output=True)
           # self.assertEqual(result.returncode, 1, result.stderr)

           # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
           _lowercase :							Optional[Any]   	=	'1'
           _lowercase :							Optional[Any]   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())
       @require_torch
       def     UpperCamelCase  (      self) ->     Optional[Any]:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	'\nfrom transformers import pipeline\n        '
           _lowercase :							Dict   	=	'\nmname = "hf-internal-testing/tiny-random-bert"\npipe = pipeline(model=mname)\n        '

           _lowercase :							Optional[Any]   	=	'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Offline mode is enabled")\nsocket.socket = offline_socket\n        '
           _lowercase :							Tuple   	=	self.get_env()
           _lowercase :							Tuple   	=	'1'
           _lowercase :							Union[str, Any]   	=	[sys.executable, '-c', '\n'.join([load, mock, run])]
           _lowercase :							Tuple   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   1,   result.stderr)
           self.assertIn(
               'You cannot infer task automatically within `pipeline` when using offline mode',   result.stderr.decode().replace('\n',   ''),   )




       @require_torch
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           _lowercase :							Dict   	=	'\nfrom transformers import AutoModel\n        '
           _lowercase :							int   	=	'\nmname = "hf-internal-testing/test_dynamic_model"\nAutoModel.from_pretrained(mname, trust_remote_code=True)\nprint("success")\n        '

           # baseline - just load from_pretrained with normal network
           _lowercase :							Optional[int]   	=	[sys.executable, '-c', '\n'.join([load, run])]

           # should succeed
           _lowercase :							int   	=	self.get_env()
           _lowercase :							List[str]   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())

           # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
           _lowercase :							Tuple   	=	'1'
           _lowercase :							Dict   	=	subprocess.run(lowerCamelCase,   env=lowerCamelCase,   check=lowerCamelCase,   capture_output=lowerCamelCase)
           self.assertEqual(result.returncode,   0,   result.stderr)
           self.assertIn('success',   result.stdout.decode())


