


import argparse
import os

# New Code #
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 import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, 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)
#
# 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
#
########################################################################


snake_case    :   Optional[int]	    =					16
snake_case    :   Any	    =					32
def     __lowerCamelCase						(			UpperCAmelCase_						:		Accelerator			,  UpperCAmelCase_						:		int = 16						):




      """simple docstring"""





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

      def tokenize_function(UpperCAmelCase_						:		Union[str, Any]						):
            # max_length=None => use the model max length (it's actually the default)
            a						:List[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():
            a						:Optional[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
      a						:Union[str, Any]										=	tokenized_datasets.rename_column('''label'''			,  '''labels'''						)

      def collate_fn(UpperCAmelCase_						:		Tuple						):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            a						:Union[str, Any]										=	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":
                  a						:Dict										=	16
            elif accelerator.mixed_precision != "no":
                  a						:Optional[int]										=	8
            else:
                  a						:Optional[Any]										=	None

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

      # Instantiate dataloaders.
      a						:Dict										=	DataLoader(
          tokenized_datasets['''train''']			,  shuffle=UpperCAmelCase_			,  collate_fn=UpperCAmelCase_			,  batch_size=UpperCAmelCase_						)
      a						:Tuple										=	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

     snake_case    :   Optional[Any]	    =					mocked_dataloaders  # noqa: F811
def     __lowerCamelCase						(			UpperCAmelCase_						:		Tuple			,  UpperCAmelCase_						:		Optional[int]						):




      """simple docstring"""





      if os.environ.get('''TESTING_MOCKED_DATALOADERS'''			,  UpperCAmelCase_						) == "1":
            a						:Any										=	2
      # Initialize accelerator
      a						:Optional[Any]										=	Accelerator(cpu=args.cpu			,  mixed_precision=args.mixed_precision						)
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      a						:Optional[Any]										=	config['''lr''']
      a						:str										=	int(config['''num_epochs''']						)
      a						:Optional[Any]										=	int(config['''seed''']						)
      a						:Union[str, Any]										=	int(config['''batch_size''']						)

      a						:Any										=	evaluate.load('''glue'''			,  '''mrpc'''						)

      # New Code #
      # We now can define an inner training loop function. It should take a batch size as the only parameter,
      # and build the dataloaders in there.
      # It also gets our decorator
      @find_executable_batch_size(starting_batch_size=UpperCAmelCase_						)
      def inner_training_loop(UpperCAmelCase_						:		List[Any]						):
            # And now just move everything below under this function
            # We need to bring in the Accelerator object from earlier
            nonlocal accelerator
            # And reset all of its attributes that could hold onto any memory:
            accelerator.free_memory()

            # Then we can declare the model, optimizer, and everything else:
            set_seed(UpperCAmelCase_						)

            # Instantiate the model (we build the model here so that the seed also control new weights initialization)
            a						:Optional[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).
            a						:str										=	model.to(accelerator.device						)

            # Instantiate optimizer
            a						:Optional[Any]										=	AdamW(params=model.parameters()			,  lr=UpperCAmelCase_						)
            a							,    a						:Optional[int]										=	get_dataloaders(UpperCAmelCase_			,  UpperCAmelCase_						)

            # Instantiate scheduler
            a						: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.
            a							,    a							,    a							,    a							,    a						:Optional[int]										=	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						)
                        a						:List[str]										=	model(**UpperCAmelCase_						)
                        a						:List[Any]										=	outputs.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():
                              a						:Any										=	model(**UpperCAmelCase_						)
                        a						:Tuple										=	outputs.logits.argmax(dim=-1						)
                        a							,    a						:Any										=	accelerator.gather_for_metrics((predictions, batch['''labels'''])						)
                        metric.add_batch(
                            predictions=UpperCAmelCase_			,  references=UpperCAmelCase_			,  )

                  a						:str										=	metric.compute()
                  # Use accelerator.print to print only on the main process.
                  accelerator.print(F'''epoch {epoch}:'''			,  UpperCAmelCase_						)

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
      inner_training_loop()
def     __lowerCamelCase						(			):




      """simple docstring"""





      a						: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.'''			,  )
      parser.add_argument('''--cpu'''			,  action='''store_true'''			,  help='''If passed, will train on the CPU.'''						)
      a						:Optional[Any]										=	parser.parse_args()
      a						:Optional[int]										=	{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
      training_function(UpperCAmelCase_			,  UpperCAmelCase_						)


if __name__ == "__main__":
     main()









import numpy as np
from matplotlib import pyplot as plt
from sklearn.datasets import load_iris
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
def     __lowerCamelCase						(			UpperCAmelCase_						:		dict						):




      """simple docstring"""





      return (data["data"], data["target"])
def     __lowerCamelCase						(			UpperCAmelCase_						:		np.ndarray			,  UpperCAmelCase_						:		np.ndarray						):




      """simple docstring"""





      a						:Optional[Any]										=	XGBClassifier()
      classifier.fit(UpperCAmelCase_			,  UpperCAmelCase_						)
      return classifier
def     __lowerCamelCase						(			):




      """simple docstring"""





      a						:List[Any]										=	load_iris()
      a							,    a						:Any										=	data_handling(UpperCAmelCase_						)
      a							,    a							,    a							,    a						:Tuple										=	train_test_split(
          UpperCAmelCase_			,  UpperCAmelCase_			,  test_size=0.25						)

      a						:List[Any]										=	iris['''target_names''']

      # Create an XGBoost Classifier from the training data
      a						:Optional[int]										=	xgboost(UpperCAmelCase_			,  UpperCAmelCase_						)

      # Display the confusion matrix of the classifier with both training and test sets
      ConfusionMatrixDisplay.from_estimator(
          UpperCAmelCase_			,  UpperCAmelCase_			,  UpperCAmelCase_			,  display_labels=UpperCAmelCase_			,  cmap='''Blues'''			,  normalize='''true'''			,  )
      plt.title('''Normalized Confusion Matrix - IRIS Dataset'''						)
      plt.show()


if __name__ == "__main__":
     import doctest

     doctest.testmod(verbose=True)
     main()





