




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


lowercase__  							=					16
lowercase__  							=					32
def      __snake_case       (   lowercase   :							Accelerator			, lowercase   :							DatasetDict			, lowercase   :							List[int]			, lowercase   :							List[int]			, lowercase   :							int = 16  ):
					snake_case_	=      AutoTokenizer.from_pretrained("bert-base-cased"  )
					snake_case_	=      DatasetDict(
					    {
					        "train": dataset["train"].select(_SCREAMING_SNAKE_CASE  ),
					        "validation": dataset["train"].select(_SCREAMING_SNAKE_CASE  ),
					        "test": dataset["validation"],
					    }  )

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

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

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

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

					snake_case_	=      DataLoader(
					    tokenized_datasets["test"]			, shuffle=_SCREAMING_SNAKE_CASE			, collate_fn=_SCREAMING_SNAKE_CASE			, batch_size=_SCREAMING_SNAKE_CASE  )

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

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

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

					set_seed(_SCREAMING_SNAKE_CASE  )

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

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

										# Instantiate optimizer
										snake_case_	=      AdamW(params=model.parameters()			, lr=_SCREAMING_SNAKE_CASE  )

										# Instantiate scheduler
										snake_case_	=      get_linear_schedule_with_warmup(
										    optimizer=_SCREAMING_SNAKE_CASE			, num_warmup_steps=100			, num_training_steps=(len(_SCREAMING_SNAKE_CASE  ) * 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.
										snake_case_	=      accelerator.prepare(
										    _SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE  )

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

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

															snake_case_	=      metric.compute()
															# Use accelerator.print to print only on the main process.
															accelerator.print(f'''epoch {epoch}:'''			, _SCREAMING_SNAKE_CASE  )

										# New Code #
										# We also run predictions on the test set at the very end
										snake_case_	=      []
										for step, batch in enumerate(_SCREAMING_SNAKE_CASE  ):
															# We could avoid this line since we set the accelerator with `device_placement=True`.
															batch.to(accelerator.device  )
															with torch.no_grad():
																				snake_case_	=      model(**_SCREAMING_SNAKE_CASE  )
															snake_case_	=      outputs.logits
															snake_case_	=      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(_SCREAMING_SNAKE_CASE			, 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:
					snake_case_	=      torch.cat(_SCREAMING_SNAKE_CASE			, dim=0  )
					snake_case_	=      torch.stack(_SCREAMING_SNAKE_CASE			, dim=0  ).sum(dim=0  ).div(int(args.num_folds  )  ).argmax(dim=-1  )
					snake_case_	=      metric.compute(predictions=_SCREAMING_SNAKE_CASE			, references=_SCREAMING_SNAKE_CASE  )
					accelerator.print("Average test metrics from all folds:"			, _SCREAMING_SNAKE_CASE  )
def      __snake_case       (   ):
					snake_case_	=      argparse.ArgumentParser(description="Simple example of training script."  )
					parser.add_argument(
					    "--mixed_precision"			, type=_SCREAMING_SNAKE_CASE			, default=_SCREAMING_SNAKE_CASE			, 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=_SCREAMING_SNAKE_CASE			, default=3			, help="The number of splits to perform across the dataset"  )
					snake_case_	=      parser.parse_args()
					snake_case_	=      {"""lr""": 2E-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
					training_function(_SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE  )


if __name__ == "__main__":
					main()






import numpy as np

import datasets


SCREAMING_SNAKE_CASE		:     Optional[int]         =	"\nCompute the Mahalanobis Distance\n\nMahalonobis distance is the distance between a point and a distribution.\nAnd not between two distinct points. It is effectively a multivariate equivalent of the Euclidean distance.\nIt was introduced by Prof. P. C. Mahalanobis in 1936\nand has been used in various statistical applications ever since\n[source: https://www.machinelearningplus.com/statistics/mahalanobis-distance/]\n"

SCREAMING_SNAKE_CASE		:     Dict         =	"\\n@article{de2000mahalanobis,\n  title={The mahalanobis distance},\n  author={De Maesschalck, Roy and Jouan-Rimbaud, Delphine and Massart, D{\'e}sir{\'e} L},\n  journal={Chemometrics and intelligent laboratory systems},\n  volume={50},\n  number={1},\n  pages={1--18},\n  year={2000},\n  publisher={Elsevier}\n}\n"

SCREAMING_SNAKE_CASE		:     List[str]         =	"\nArgs:\n    X: List of datapoints to be compared with the `reference_distribution`.\n    reference_distribution: List of datapoints from the reference distribution we want to compare to.\nReturns:\n    mahalanobis: The Mahalonobis distance for each datapoint in `X`.\nExamples:\n\n    >>> mahalanobis_metric = datasets.load_metric(\"mahalanobis\")\n    >>> results = mahalanobis_metric.compute(reference_distribution=[[0, 1], [1, 0]], X=[[0, 1]])\n    >>> print(results)\n    {'mahalanobis': array([0.5])}\n"

@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION	,	_KWARGS_DESCRIPTION				)
class        UpperCamelCase						(				datasets.Metric				):





						def        A_							(self						)				->				str:
												return datasets.MetricInfo(
												    description=_DESCRIPTION			,     citation=_CITATION			,     inputs_description=_KWARGS_DESCRIPTION			,     features=datasets.Features(
												        {
												            """X""": datasets.Sequence(datasets.Value("""float"""			,     id="""sequence"""						)			,     id="""X"""						),
												        }						)			,     )





						def        A_							(self			,     __UpperCamelCase			,     __UpperCamelCase						)				->				int:
												# convert to numpy arrays
												UpperCamelCase_		:							int  =		np.array(__UpperCamelCase						)
												UpperCamelCase_		:							Union[str, Any]  =		np.array(__UpperCamelCase						)

												# Assert that arrays are 2D
												if len(X.shape						) != 2:
																		raise ValueError("""Expected `X` to be a 2D vector"""						)
												if len(reference_distribution.shape						) != 2:
																		raise ValueError("""Expected `reference_distribution` to be a 2D vector"""						)
												if reference_distribution.shape[0] < 2:
																		raise ValueError(
																		    """Expected `reference_distribution` to be a 2D vector with more than one element in the first dimension"""						)

												# Get mahalanobis distance for each prediction
												UpperCamelCase_		:							str  =		X - np.mean(__UpperCamelCase						)
												UpperCamelCase_		:							Dict  =		np.cov(reference_distribution.T						)
												try:
																		UpperCamelCase_		:							str  =		np.linalg.inv(__UpperCamelCase						)
												except np.linalg.LinAlgError:
																		UpperCamelCase_		:							List[Any]  =		np.linalg.pinv(__UpperCamelCase						)
												UpperCamelCase_		:							Tuple  =		np.dot(__UpperCamelCase			,     __UpperCamelCase						)
												UpperCamelCase_		:							Tuple  =		np.dot(__UpperCamelCase			,     X_minus_mu.T						).diagonal()

												return {"mahalanobis": mahal_dist}






