




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


_UpperCAmelCase	:						List[Any]			=     16
_UpperCAmelCase	:						Optional[int]			=     32



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase,			lowerCamelCase,			lowerCamelCase = 16   ):
  lowercase    :List[str]      	=						AutoTokenizer.from_pretrained("bert-base-cased"   )
  lowercase    :Optional[Any]      	=						DatasetDict(
      {
          "train": dataset["train"].select(lowerCamelCase   ),
          "validation": dataset["train"].select(lowerCamelCase   ),
          "test": dataset["validation"],
      }   )

  def tokenize_function(lowerCamelCase   ):
    # max_length=None => use the model max length (it's actually the default)
    lowercase    :List[str]      	=						tokenizer(examples["sentence1"],			examples["sentence2"],			truncation=lowerCamelCase,			max_length=lowerCamelCase   )
    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():
    lowercase    :Optional[int]      	=						datasets.map(
        lowerCamelCase,			batched=lowerCamelCase,			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
  lowercase    :List[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.
    lowercase    :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":
      lowercase    :Dict      	=						16
    elif accelerator.mixed_precision != "no":
      lowercase    :Optional[int]      	=						8
    else:
      lowercase    :int      	=						None

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

  # Instantiate dataloaders.
  lowercase    :Tuple      	=						DataLoader(
      tokenized_datasets["train"],			shuffle=lowerCamelCase,			collate_fn=lowerCamelCase,			batch_size=lowerCamelCase   )
  lowercase    :Dict      	=						DataLoader(
      tokenized_datasets["validation"],			shuffle=lowerCamelCase,			collate_fn=lowerCamelCase,			batch_size=lowerCamelCase   )

  lowercase    :List[Any]      	=						DataLoader(
      tokenized_datasets["test"],			shuffle=lowerCamelCase,			collate_fn=lowerCamelCase,			batch_size=lowerCamelCase   )

  return train_dataloader, eval_dataloader, test_dataloader



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

  lowercase    :int      	=						evaluate.load("glue",			"mrpc"   )

  # If the batch size is too big we use gradient accumulation
  lowercase    :Union[str, Any]      	=						1
  if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
    lowercase    :Any      	=						batch_size // MAX_GPU_BATCH_SIZE
    lowercase    :str      	=						MAX_GPU_BATCH_SIZE

  set_seed(lowerCamelCase   )

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

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

    # Instantiate optimizer
    lowercase    :Any      	=						AdamW(params=model.parameters(),			lr=lowerCamelCase   )

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

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

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

      lowercase    :Union[str, Any]      	=						metric.compute()
      # Use accelerator.print to print only on the main process.
      accelerator.print(F"epoch {epoch}:",			lowerCamelCase   )

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



def       UpperCAmelCase__  (     ):
  lowercase    :List[str]      	=						argparse.ArgumentParser(description="Simple example of training script."   )
  parser.add_argument(
      "--mixed_precision",			type=lowerCamelCase,			default=lowerCamelCase,			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=lowerCamelCase,			default=3,			help="The number of splits to perform across the dataset"   )
  lowercase    :Union[str, Any]      	=						parser.parse_args()
  lowercase    :Dict      	=						{"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
  training_function(lowerCamelCase,			lowerCamelCase   )


if __name__ == "__main__":
     main()









import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


_UpperCAmelCase	:						List[Any]			=     logging.get_logger(__name__)



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :int      	=						r"\w+[.]\d+"
  lowercase    :Tuple      	=						re.findall(lowerCamelCase,			lowerCamelCase   )
  for pat in pats:
    lowercase    :List[str]      	=						key.replace(lowerCamelCase,			"_".join(pat.split("."   )   )   )
  return key



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase   ):
  lowercase    :List[str]      	=						pt_tuple_key[:-1] + ("scale",)
  if (
      any("norm" in str_ for str_ in pt_tuple_key   )
      and (pt_tuple_key[-1] == "bias")
      and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
      and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
  ):
    lowercase    :Optional[int]      	=						pt_tuple_key[:-1] + ("scale",)
    return renamed_pt_tuple_key, pt_tensor
  elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
    lowercase    :int      	=						pt_tuple_key[:-1] + ("scale",)
    return renamed_pt_tuple_key, pt_tensor

  # embedding
  if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
    lowercase    :List[str]      	=						pt_tuple_key[:-1] + ("embedding",)
    return renamed_pt_tuple_key, pt_tensor

  # conv layer
  lowercase    :Union[str, Any]      	=						pt_tuple_key[:-1] + ("kernel",)
  if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
    lowercase    :List[Any]      	=						pt_tensor.transpose(2,			3,			1,			0   )
    return renamed_pt_tuple_key, pt_tensor

  # linear layer
  lowercase    :Optional[Any]      	=						pt_tuple_key[:-1] + ("kernel",)
  if pt_tuple_key[-1] == "weight":
    lowercase    :Tuple      	=						pt_tensor.T
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm weight
  lowercase    :Tuple      	=						pt_tuple_key[:-1] + ("weight",)
  if pt_tuple_key[-1] == "gamma":
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm bias
  lowercase    :int      	=						pt_tuple_key[:-1] + ("bias",)
  if pt_tuple_key[-1] == "beta":
    return renamed_pt_tuple_key, pt_tensor

  return pt_tuple_key, pt_tensor







def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase=42   ):
  # Step 1: Convert pytorch tensor to numpy
  lowercase    :Optional[int]      	=						{k: v.numpy() for k, v in pt_state_dict.items()}

  # Step 2: Since the model is stateless, get random Flax params
  lowercase    :str      	=						flax_model.init_weights(PRNGKey(lowerCamelCase   )   )

  lowercase    :Tuple      	=						flatten_dict(lowerCamelCase   )
  lowercase    :Optional[Any]      	=						{}

  # Need to change some parameters name to match Flax names
  for pt_key, pt_tensor in pt_state_dict.items():
    lowercase    :List[Any]      	=						rename_key(lowerCamelCase   )
    lowercase    :List[Any]      	=						tuple(renamed_pt_key.split("."   )   )

    # Correctly rename weight parameters
    lowercase	,    lowercase    :List[str]      	=						rename_key_and_reshape_tensor(lowerCamelCase,			lowerCamelCase,			lowerCamelCase   )

    if flax_key in random_flax_state_dict:
      if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
        raise ValueError(
            F"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "
            F"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."   )

        # also add unexpected weight so that warning is thrown
    lowercase    :List[str]      	=						jnp.asarray(lowerCamelCase   )

  return unflatten_dict(lowerCamelCase   )



