
"""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 					_UpperCAmelCase      ( __lowerCamelCase   :       int ,					__lowerCamelCase   :       str = 16		)     ->					List[Any]:
  _snake_case					     =       AutoTokenizer.from_pretrained('''bert-base-cased'''		)
  _snake_case					     =       load_dataset('''glue''' ,					'''mrpc'''		)

  def tokenize_function(__lowerCamelCase   :       Tuple		):
    # max_length=None => use the model max length (it's actually the default)
    _snake_case					     =       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():
    _snake_case					     =       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
  _snake_case					     =       tokenized_datasets.rename_column('''label''' ,					'''labels'''		)

  def collate_fn(__lowerCamelCase   :       Optional[Any]		):
    # On TPU it's best to pad everything to the same length or training will be very slow.
    _snake_case					     =       1_28 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(
        UpperCamelCase__ ,					padding='''longest''' ,					max_length=UpperCamelCase__ ,					pad_to_multiple_of=UpperCamelCase__ ,					return_tensors='''pt''' ,					)

  # Instantiate dataloaders.
  _snake_case					     =       DataLoader(
      tokenized_datasets['''train'''] ,					shuffle=UpperCamelCase__ ,					collate_fn=UpperCamelCase__ ,					batch_size=UpperCamelCase__		)
  _snake_case					     =       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

 UpperCAmelCase__          = mocked_dataloaders  # noqa: F811
def 					_UpperCAmelCase      ( __lowerCamelCase   :       Optional[Any] ,					__lowerCamelCase   :       Dict		)     ->					Dict:





  if os.environ.get('''TESTING_MOCKED_DATALOADERS''' ,					UpperCamelCase__		) == "1":
    _snake_case					     =       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:
    _snake_case					     =       Accelerator(
        cpu=args.cpu ,					mixed_precision=args.mixed_precision ,					log_with='''all''' ,					project_dir=args.project_dir		)
  else:
    _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''']		)
  set_seed(UpperCamelCase__		)

  _snake_case     ,					_snake_case					     =       get_dataloaders(UpperCamelCase__ ,					UpperCamelCase__		)
  _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

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

  # Instantiate optimizer
  _snake_case					     =       AdamW(params=model.parameters() ,					lr=UpperCamelCase__		)

  # Instantiate scheduler
  _snake_case					     =       get_linear_schedule_with_warmup(
      optimizer=UpperCamelCase__ ,					num_warmup_steps=1_00 ,					num_training_steps=(len(UpperCamelCase__		) * 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     ,					_snake_case     ,					_snake_case     ,					_snake_case     ,					_snake_case					     =       accelerator.prepare(
      UpperCamelCase__ ,					UpperCamelCase__ ,					UpperCamelCase__ ,					UpperCamelCase__ ,					UpperCamelCase__		)

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

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

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

    # 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(UpperCamelCase__		),
              '''epoch''': epoch,
          } ,					step=UpperCamelCase__ ,					)

    # 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 					_UpperCAmelCase      ( )     ->					Dict:
  _snake_case					     =       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.'''		)
  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=UpperCamelCase__ ,					default='''logs''' ,					help='''Location on where to store experiment tracking logs` and relevent project information''' ,					)
  _snake_case					     =       parser.parse_args()
  _snake_case					     =       {'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
  training_function(UpperCamelCase__ ,					UpperCamelCase__		)


if __name__ == "__main__":
 main()




"""simple docstring"""




from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase							 =     logging.get_logger(__name__)

__lowerCamelCase							 =     {
    "huggingface/time-series-transformer-tourism-monthly": (
        "https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}
class 			UpperCamelCase__(	__A       ):
						lowerCAmelCase__     :  List[Any]  							=						'time_series_transformer'
						lowerCAmelCase__     :  Dict  							=						{
						    'hidden_size': 'd_model',
						    'num_attention_heads': 'encoder_attention_heads',
						    'num_hidden_layers': 'encoder_layers',
						}





						def __init__(				self						,__UpperCAmelCase = None						,__UpperCAmelCase = None						,__UpperCAmelCase = "student_t"						,__UpperCAmelCase = "nll"						,__UpperCAmelCase = 1						,__UpperCAmelCase = [1, 2, 3, 4, 5, 6, 7]						,__UpperCAmelCase = "mean"						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = None						,__UpperCAmelCase = None						,__UpperCAmelCase = 32						,__UpperCAmelCase = 32						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = True						,__UpperCAmelCase = "gelu"						,__UpperCAmelCase = 64						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 1_00						,__UpperCAmelCase = 0.0_2						,__UpperCAmelCase=True						,**__UpperCAmelCase						,)    -> List[str]:
													# time series specific configuration
													A__		   =							prediction_length
													A__		   =							context_length or prediction_length
													A__		   =							distribution_output
													A__		   =							loss
													A__		   =							input_size
													A__		   =							num_time_features
													A__		   =							lags_sequence
													A__		   =							scaling
													A__		   =							num_dynamic_real_features
													A__		   =							num_static_real_features
													A__		   =							num_static_categorical_features
													if cardinality and num_static_categorical_features > 0:
																				if len(__UpperCAmelCase   ) != num_static_categorical_features:
																											raise ValueError(
																											    'The cardinality should be a list of the same length as `num_static_categorical_features`'   )
																				A__		   =							cardinality
													else:
																				A__		   =							[0]
													if embedding_dimension and num_static_categorical_features > 0:
																				if len(__UpperCAmelCase   ) != num_static_categorical_features:
																											raise ValueError(
																											    'The embedding dimension should be a list of the same length as `num_static_categorical_features`'   )
																				A__		   =							embedding_dimension
													else:
																				A__		   =							[min(50						,(cat + 1) // 2   ) for cat in self.cardinality]
													A__		   =							num_parallel_samples

													# Transformer architecture configuration
													A__		   =							input_size * len(__UpperCAmelCase   ) + self._number_of_features
													A__		   =							d_model
													A__		   =							encoder_attention_heads
													A__		   =							decoder_attention_heads
													A__		   =							encoder_ffn_dim
													A__		   =							decoder_ffn_dim
													A__		   =							encoder_layers
													A__		   =							decoder_layers

													A__		   =							dropout
													A__		   =							attention_dropout
													A__		   =							activation_dropout
													A__		   =							encoder_layerdrop
													A__		   =							decoder_layerdrop

													A__		   =							activation_function
													A__		   =							init_std

													A__		   =							use_cache

													super().__init__(is_encoder_decoder=__UpperCAmelCase						,**__UpperCAmelCase   )





						@property
						def  snake_case__		(				self   )    -> int:
													return (
													    sum(self.embedding_dimension   )
													    + self.num_dynamic_real_features
													    + self.num_time_features
													    + self.num_static_real_features
													    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
													)





