



'''simple docstring'''





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,
# specifically showcasing the experiment tracking capability,
# 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											=		16
_UpperCamelCase											=		32

def 			a_	(				_lowerCAmelCase			,_lowerCAmelCase = 16			) -> Any:
  __lowerCamelCase			:			List[Any]       			=					AutoTokenizer.from_pretrained('bert-base-cased'			)
  __lowerCamelCase			:			List[str]       			=					load_dataset('glue'			,'mrpc'			)

  def tokenize_function(_lowerCAmelCase			):
    # max_length=None => use the model max length (it's actually the default)
    __lowerCamelCase			:			Optional[int]       			=					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():
    __lowerCamelCase			:			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
  __lowerCamelCase			:			Dict       			=					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.
    __lowerCamelCase			:			Optional[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":
      __lowerCamelCase			:			Optional[Any]       			=					16
    elif accelerator.mixed_precision != "no":
      __lowerCamelCase			:			List[str]       			=					8
    else:
      __lowerCamelCase			:			Optional[int]       			=					None

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

  # Instantiate dataloaders.
  __lowerCamelCase			:			List[Any]       			=					DataLoader(
      tokenized_datasets['train']			,shuffle=__lowerCAmelCase			,collate_fn=__lowerCAmelCase			,batch_size=__lowerCAmelCase			)
  __lowerCamelCase			:			List[str]       			=					DataLoader(
      tokenized_datasets['validation']			,shuffle=__lowerCAmelCase			,collate_fn=__lowerCAmelCase			,batch_size=__lowerCAmelCase			)

  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

   _UpperCamelCase											=		mocked_dataloaders  # noqa: F811

def 			a_	(				_lowerCAmelCase			,_lowerCAmelCase			) -> List[Any]:
  # For testing only
  if os.environ.get('TESTING_MOCKED_DATALOADERS'			,__lowerCAmelCase			) == "1":
    __lowerCamelCase			:			Any       			=					2
  # Initialize Accelerator

  # New Code #
  # We pass in "all" to `log_with` to grab all available trackers in the environment
  # Note: If using a custom `Tracker` class, should be passed in here such as:
  # >>> log_with = ["all", MyCustomTrackerClassInstance()]
  if args.with_tracking:
    __lowerCamelCase			:			Dict       			=					Accelerator(
        cpu=args.cpu			,mixed_precision=args.mixed_precision			,log_with='all'			,project_dir=args.project_dir			)
  else:
    __lowerCamelCase			:			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
  __lowerCamelCase			:			Union[str, Any]       			=					config['lr']
  __lowerCamelCase			:			Optional[int]       			=					int(config['num_epochs']			)
  __lowerCamelCase			:			Optional[int]       			=					int(config['seed']			)
  __lowerCamelCase			:			List[Any]       			=					int(config['batch_size']			)
  set_seed(__lowerCAmelCase			)

  __lowerCamelCase			,__lowerCamelCase			:			Any       			=					get_dataloaders(__lowerCAmelCase			,__lowerCAmelCase			)
  __lowerCamelCase			:			Tuple       			=					evaluate.load('glue'			,'mrpc'			)

  # If the batch size is too big we use gradient accumulation
  __lowerCamelCase			:			Dict       			=					1
  if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
    __lowerCamelCase			:			Tuple       			=					batch_size // MAX_GPU_BATCH_SIZE
    __lowerCamelCase			:			Optional[int]       			=					MAX_GPU_BATCH_SIZE

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

  # Instantiate optimizer
  __lowerCamelCase			:			str       			=					AdamW(params=model.parameters()			,lr=__lowerCAmelCase			)

  # Instantiate scheduler
  __lowerCamelCase			:			List[str]       			=					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.
  __lowerCamelCase			,__lowerCamelCase			,__lowerCamelCase			,__lowerCamelCase			,__lowerCamelCase			:			Union[str, Any]       			=					accelerator.prepare(
      __lowerCAmelCase			,__lowerCAmelCase			,__lowerCAmelCase			,__lowerCAmelCase			,__lowerCAmelCase			)

  # New Code #
  # We need to initialize the trackers we use. Overall configurations can also be stored
  if args.with_tracking:
    __lowerCamelCase			:			int       			=					os.path.split(__lowerCAmelCase			)[-1].split('.'			)[0]
    accelerator.init_trackers(__lowerCAmelCase			,__lowerCAmelCase			)

  # Now we train the model
  for epoch in range(__lowerCAmelCase			):
    model.train()
    # New Code #
    # For our tracking example, we will log the total loss of each epoch
    if args.with_tracking:
      __lowerCamelCase			:			Dict       			=					0
    for step, batch in enumerate(__lowerCAmelCase			):
      # We could avoid this line since we set the accelerator with `device_placement=True`.
      batch.to(accelerator.device			)
      __lowerCamelCase			:			Union[str, Any]       			=					model(**__lowerCAmelCase			)
      __lowerCamelCase			:			Tuple       			=					outputs.loss
      # New Code #
      if args.with_tracking:
        total_loss += loss.detach().float()
      __lowerCamelCase			:			Tuple       			=					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` (the default).
      batch.to(accelerator.device			)
      with torch.no_grad():
        __lowerCamelCase			:			List[Any]       			=					model(**__lowerCAmelCase			)
      __lowerCamelCase			:			int       			=					outputs.logits.argmax(dim=-1			)
      __lowerCamelCase			,__lowerCamelCase			:			Any       			=					accelerator.gather_for_metrics((predictions, batch['labels'])			)
      metric.add_batch(
          predictions=__lowerCAmelCase			,references=__lowerCAmelCase			,)

    __lowerCamelCase			:			Dict       			=					metric.compute()
    # Use accelerator.print to print only on the main process.
    accelerator.print(F'epoch {epoch}:'			,__lowerCAmelCase			)

    # New Code #
    # To actually log, we call `Accelerator.log`
    # The values passed can be of `str`, `int`, `float` or `dict` of `str` to `float`/`int`
    if args.with_tracking:
      accelerator.log(
          {
              'accuracy': eval_metric['accuracy'],
              'f1': eval_metric['f1'],
              'train_loss': total_loss.item() / len(__lowerCAmelCase			),
              'epoch': epoch,
          }			,step=__lowerCAmelCase			,)

    # New Code #
    # When a run is finished, you should call `accelerator.end_training()`
    # to close all of the open trackers
  if args.with_tracking:
    accelerator.end_training()

def 			a_	(				) -> int:
  __lowerCamelCase			:			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.'			)
  parser.add_argument(
      '--with_tracking'			,action='store_true'			,help='Whether to load in all available experiment trackers from the environment and use them for logging.'			,)
  parser.add_argument(
      '--project_dir'			,type=__lowerCAmelCase			,default='logs'			,help='Location on where to store experiment tracking logs` and relevent project information'			,)
  __lowerCamelCase			:			List[str]       			=					parser.parse_args()
  __lowerCamelCase			:			str       			=					{'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
  training_function(__lowerCAmelCase			,__lowerCAmelCase			)


if __name__ == "__main__":
   main()












import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
__magic_name__  :  Tuple					=       data_utils.TransfoXLTokenizer
__magic_name__  :  int					=       data_utils.TransfoXLCorpus
__magic_name__  :  List[str]					=       data_utils
__magic_name__  :  str					=       data_utils
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase				):
  if transfo_xl_dataset_file:
    # Convert a pre-processed corpus (see original TensorFlow repo)
    with open(__lowerCAmelCase   , '''rb'''				) as fp:
      lowerCAmelCase__										=      pickle.load(__lowerCAmelCase   , encoding='''latin1'''				)
    # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
    lowerCAmelCase__										=      pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
    print(F"""Save vocabulary to {pytorch_vocab_dump_path}"""				)
    lowerCAmelCase__										=      corpus.vocab.__dict__
    torch.save(__lowerCAmelCase   , __lowerCAmelCase				)

    lowerCAmelCase__										=      corpus.__dict__
    corpus_dict_no_vocab.pop('''vocab'''   , __lowerCAmelCase				)
    lowerCAmelCase__										=      pytorch_dump_folder_path + '''/''' + CORPUS_NAME
    print(F"""Save dataset to {pytorch_dataset_dump_path}"""				)
    torch.save(__lowerCAmelCase   , __lowerCAmelCase				)

  if tf_checkpoint_path:
    # Convert a pre-trained TensorFlow model
    lowerCAmelCase__										=      os.path.abspath(__lowerCAmelCase				)
    lowerCAmelCase__										=      os.path.abspath(__lowerCAmelCase				)

    print(F"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""				)
    # Initialise PyTorch model
    if transfo_xl_config_file == "":
      lowerCAmelCase__										=      TransfoXLConfig()
    else:
      lowerCAmelCase__										=      TransfoXLConfig.from_json_file(__lowerCAmelCase				)
    print(F"""Building PyTorch model from configuration: {config}"""				)
    lowerCAmelCase__										=      TransfoXLLMHeadModel(__lowerCAmelCase				)

    lowerCAmelCase__										=      load_tf_weights_in_transfo_xl(__lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase				)
    # Save pytorch-model
    lowerCAmelCase__										=      os.path.join(__lowerCAmelCase   , __lowerCAmelCase				)
    lowerCAmelCase__										=      os.path.join(__lowerCAmelCase   , __lowerCAmelCase				)
    print(F"""Save PyTorch model to {os.path.abspath(__lowerCAmelCase				)}"""				)
    torch.save(model.state_dict()   , __lowerCAmelCase				)
    print(F"""Save configuration file to {os.path.abspath(__lowerCAmelCase				)}"""				)
    with open(__lowerCAmelCase   , '''w'''   , encoding='''utf-8'''				) as f:
      f.write(config.to_json_string()				)


if __name__ == "__main__":
      __magic_name__  :  int					=       argparse.ArgumentParser()
      parser.add_argument(
          """--pytorch_dump_folder_path""",
          default=None,
          type=str,
          required=True,
          help="""Path to the folder to store the PyTorch model or dataset/vocab.""",
      )
      parser.add_argument(
          """--tf_checkpoint_path""",
          default="""""",
          type=str,
          help="""An optional path to a TensorFlow checkpoint path to be converted.""",
      )
      parser.add_argument(
          """--transfo_xl_config_file""",
          default="""""",
          type=str,
          help=(
              """An optional config json file corresponding to the pre-trained BERT model. \n"""
              """This specifies the model architecture."""
          ),
      )
      parser.add_argument(
          """--transfo_xl_dataset_file""",
          default="""""",
          type=str,
          help="""An optional dataset file to be converted in a vocabulary.""",
      )
      __magic_name__  :  Optional[Any]					=       parser.parse_args()
      convert_transfo_xl_checkpoint_to_pytorch(
          args.tf_checkpoint_path,
          args.transfo_xl_config_file,
          args.pytorch_dump_folder_path,
          args.transfo_xl_dataset_file,
      )

