code: def main():
      args = parse_args()
  
      # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
      # If we're using tracking, we also need to initialize it here and it will pick up all supported trackers in the environment
      accelerator = (
          Accelerator(log_with="all", logging_dir=args.output_dir)
          if args.with_tracking
          else Accelerator()
      )
      # Make one log on every process with the configuration for debugging.
      logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO,
      )
      logger.info(accelerator.state)
  
      # Setup logging, we only want one process per machine to log things on the screen.
      # accelerator.is_local_main_process is only True for one process per machine.
      logger.setLevel(
          logging.INFO
          if accelerator.is_local_main_process
          else logging.ERROR
      )
      if accelerator.is_local_main_process:
          datasets.utils.logging.set_verbosity_warning()
          transformers.utils.logging.set_verbosity_info()
      else:
          datasets.utils.logging.set_verbosity_error()
          transformers.utils.logging.set_verbosity_error()
  
      # If passed along, set the training seed now.
      if args.seed is not None:
          set_seed(args.seed)
  
      # Handle the repository creation
      if accelerator.is_main_process:
          if args.push_to_hub:
              if args.hub_model_id is None:
                  repo_name = get_full_repo_name(
                      Path(args.output_dir).name, token=args.hub_token
                  )
              else:
                  repo_name = args.hub_model_id
              repo = Repository(args.output_dir, clone_from=repo_name)
  
              with open(
                  os.path.join(args.output_dir, ".gitignore"), "w+"
              ) as gitignore:
                  if "step_*" not in gitignore:
                      gitignore.write("step_*\n")
                  if "epoch_*" not in gitignore:
                      gitignore.write("epoch_*\n")
          elif args.output_dir is not None:
              os.makedirs(args.output_dir, exist_ok=True)
      accelerator.wait_for_everyone()
  
      # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
      # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
      # (the dataset will be downloaded automatically from the datasets Hub).
      #
      # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
      # 'text' is found. You can easily tweak this behavior (see below).
      #
      # In distributed training, the load_dataset function guarantee that only one local process can concurrently
      # download the dataset.
      if args.dataset_name is not None:
          # Downloading and loading a dataset from the hub.
          raw_datasets = datasets.DatasetDict(
              {
                  "train": datasets.Dataset.from_dict(
                      load_dataset(
                          args.dataset_name, args.dataset_config_name
                      )["train"][: args.n_train + args.n_val]
                  )
              }
          )
          if "validation" not in raw_datasets.keys():
              raw_datasets["validation"] = load_dataset(
                  args.dataset_name,
                  args.dataset_config_name,
                  split=f"train[:{args.validation_split_percentage}%]",
              )
              raw_datasets["train"] = load_dataset(
                  args.dataset_name,
                  args.dataset_config_name,
                  split=f"train[{args.validation_split_percentage}%:]",
              )
      else:
          data_files = {}
          dataset_args = {}
          if args.train_file is not None:
              data_files["train"] = args.train_file
          if args.validation_file is not None:
              data_files["validation"] = args.validation_file
          extension = args.train_file.split(".")[-1]
          if extension == "txt":
              extension = "text"
              dataset_args[
                  "keep_linebreaks"
              ] = not args.no_keep_linebreaks
          raw_datasets = load_dataset(
              extension, data_files=data_files, **dataset_args
          )
          # If no validation data is there, validation_split_percentage will be used to divide the dataset.
          if "validation" not in raw_datasets.keys():
              raw_datasets["validation"] = load_dataset(
                  extension,
                  data_files=data_files,
                  split=f"train[:{args.validation_split_percentage}%]",
                  **dataset_args,
              )
              raw_datasets["train"] = load_dataset(
                  extension,
                  data_files=data_files,
                  split=f"train[{args.validation_split_percentage}%:]",
                  **dataset_args,
              )
  
      # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
      # https://huggingface.co/docs/datasets/loading_datasets.html.
  
      # Load pretrained model and tokenizer
      #
      # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
      # download model & vocab.
      if args.config_name:
          config = AutoConfig.from_pretrained(args.config_name)
      elif args.model_name_or_path:
          config = AutoConfig.from_pretrained(args.model_name_or_path)
      else:
          config = CONFIG_MAPPING[args.model_type]()
          logger.warning(
              "You are instantiating a new config instance from scratch."
          )
  
      if args.tokenizer_name:
          tokenizer = AutoTokenizer.from_pretrained(
              args.tokenizer_name, use_fast=not args.use_slow_tokenizer
          )
      elif args.model_name_or_path:
          tokenizer = AutoTokenizer.from_pretrained(
              args.model_name_or_path,
              use_fast=not args.use_slow_tokenizer,
          )
      else:
          raise ValueError(
              "You are instantiating a new tokenizer from scratch. This is not supported by this script."
              "You can do it from another script, save it, and load it from here, using --tokenizer_name."
          )
  
      if args.model_name_or_path:
          model = AutoModelForCausalLM.from_pretrained(
              args.model_name_or_path,
              from_tf=bool(".ckpt" in args.model_name_or_path),
              config=config,
          )
      else:
          logger.info("Training new model from scratch")
          model = AutoModelForCausalLM.from_config(config)
  
      model.resize_token_embeddings(len(tokenizer))
  
      # Preprocessing the datasets.
      # First we tokenize all the texts.
      column_names = raw_datasets["train"].column_names
      text_column_name = (
          "text" if "text" in column_names else column_names[0]
      )
  
      def tokenize_function(examples):
          return tokenizer(examples[text_column_name])
  
      with accelerator.main_process_first():
          tokenized_datasets = raw_datasets.map(
              tokenize_function,
              batched=True,
              num_proc=args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not args.overwrite_cache,
              desc="Running tokenizer on dataset",
          )
  
      if args.block_size is None:
          block_size = tokenizer.model_max_length
          if block_size > 1024:
              logger.warning(
                  f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
                  "Picking 1024 instead. You can change that default value by passing --block_size xxx."
              )
          block_size = 1024
      else:
          if args.block_size > tokenizer.model_max_length:
              logger.warning(
                  f"The block_size passed ({args.block_size}) is larger than the maximum length for the model"
                  f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
              )
          block_size = min(args.block_size, tokenizer.model_max_length)
  
      # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
      def group_texts(examples):
          # Concatenate all texts.
          concatenated_examples = {
              k: list(chain(*examples[k])) for k in examples.keys()
          }
          total_length = len(
              concatenated_examples[list(examples.keys())[0]]
          )
          # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
          # customize this part to your needs.
          if total_length >= block_size:
              total_length = (total_length // block_size) * block_size
          # Split by chunks of max_len.
          result = {
              k: [
                  t[i : i + block_size]
                  for i in range(0, total_length, block_size)
              ]
              for k, t in concatenated_examples.items()
          }
          result["labels"] = result["input_ids"].copy()
          return result
  
      # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
      # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
      # to preprocess.
      #
      # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
      # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
  
      with accelerator.main_process_first():
          lm_datasets = tokenized_datasets.map(
              group_texts,
              batched=True,
              num_proc=args.preprocessing_num_workers,
              load_from_cache_file=not args.overwrite_cache,
              desc=f"Grouping texts in chunks of {block_size}",
          )
  
      train_dataset = lm_datasets["train"]
      eval_dataset = lm_datasets["validation"]
  
      # Log a few random samples from the training set:
      for index in random.sample(range(len(train_dataset)), 3):
          logger.info(
              f"Sample {index} of the training set: {train_dataset[index]}."
          )
  
      # DataLoaders creation:
      train_dataloader = DataLoader(
          train_dataset,
          shuffle=True,
          collate_fn=default_data_collator,
          batch_size=args.per_device_train_batch_size,
      )
      eval_dataloader = DataLoader(
          eval_dataset,
          collate_fn=default_data_collator,
          batch_size=args.per_device_eval_batch_size,
      )
  
      # Optimizer
      # Split weights in two groups, one with weight decay and the other not.
      no_decay = ["bias", "LayerNorm.weight"]
      optimizer_grouped_parameters = [
          {
              "params": [
                  p
                  for n, p in model.named_parameters()
                  if not any(nd in n for nd in no_decay)
              ],
              "weight_decay": args.weight_decay,
          },
          {
              "params": [
                  p
                  for n, p in model.named_parameters()
                  if any(nd in n for nd in no_decay)
              ],
              "weight_decay": 0.0,
          },
      ]
      optimizer = AdamW(
          optimizer_grouped_parameters, lr=args.learning_rate
      )
  
      # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.
      if accelerator.distributed_type == DistributedType.TPU:
          model.tie_weights()
  
      # Scheduler and math around the number of training steps.
      num_update_steps_per_epoch = math.ceil(
          len(train_dataloader) / args.gradient_accumulation_steps
      )
      if args.max_train_steps is None:
          args.max_train_steps = (
              args.num_train_epochs * num_update_steps_per_epoch
          )
      else:
          args.num_train_epochs = math.ceil(
              args.max_train_steps / num_update_steps_per_epoch
          )
  
      lr_scheduler = get_scheduler(
          name=args.lr_scheduler_type,
          optimizer=optimizer,
          num_warmup_steps=args.num_warmup_steps,
          num_training_steps=args.max_train_steps,
      )
  
      # Prepare everything with our `accelerator`.
      (
          model,
          optimizer,
          train_dataloader,
          eval_dataloader,
          lr_scheduler,
      ) = accelerator.prepare(
          model,
          optimizer,
          train_dataloader,
          eval_dataloader,
          lr_scheduler,
      )
  
      # Figure out how many steps we should save the Accelerator states
      if hasattr(args.checkpointing_steps, "isdigit"):
          checkpointing_steps = args.checkpointing_steps
          if args.checkpointing_steps.isdigit():
              checkpointing_steps = int(args.checkpointing_steps)
      else:
          checkpointing_steps = None
  
      # We need to initialize the trackers we use, and also store our configuration
      if args.with_tracking:
          experiment_config = vars(args)
          # TensorBoard cannot log Enums, need the raw value
          experiment_config["lr_scheduler_type"] = experiment_config[
              "lr_scheduler_type"
          ].value
          accelerator.init_trackers("clm_no_trainer", experiment_config)
  
      # Train!
      total_batch_size = (
          args.per_device_train_batch_size
          * accelerator.num_processes
          * args.gradient_accumulation_steps
      )
  
      logger.info("***** Running training *****")
      logger.info(f"  Num examples = {len(train_dataset)}")
      logger.info(f"  Num Epochs = {args.num_train_epochs}")
      logger.info(
          f"  Instantaneous batch size per device = {args.per_device_train_batch_size}"
      )
      logger.info(
          f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
      )
      logger.info(
          f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}"
      )
      logger.info(
          f"  Total optimization steps = {int(args.max_train_steps/accelerator.num_processes)}"
      )
      # Only show the progress bar once on each machine.
      progress_bar = tqdm(
          range(int(args.max_train_steps / accelerator.num_processes)),
          disable=not accelerator.is_local_main_process,
      )
      completed_steps = 0
  
      # Potentially load in the weights and states from a previous save
      if args.resume_from_checkpoint:
          if (
              args.resume_from_checkpoint is not None
              or args.resume_from_checkpoint != ""
          ):
              accelerator.print(
                  f"Resumed from checkpoint: {args.resume_from_checkpoint}"
              )
              accelerator.load_state(args.resume_from_checkpoint)
              resume_step = None
              path = args.resume_from_checkpoint
          else:
              # Get the most recent checkpoint
              dirs = [
                  f.name for f in os.scandir(os.getcwd()) if f.is_dir()
              ]
              dirs.sort(key=os.path.getctime)
              path = dirs[
                  -1
              ]  # Sorts folders by date modified, most recent checkpoint is the last
          if "epoch" in path:
              args.num_train_epochs -= int(path.replace("epoch_", ""))
          else:
              resume_step = int(path.replace("step_", ""))
              args.num_train_epochs -= resume_step // len(
                  train_dataloader
              )
              resume_step = (
                  args.num_train_epochs * len(train_dataloader)
              ) - resume_step
  
      for epoch in range(args.num_train_epochs):
          model.train()
          if args.with_tracking:
              total_loss = 0
          for step, batch in enumerate(train_dataloader):
              # We need to skip steps until we reach the resumed step
              if (
                  args.resume_from_checkpoint
                  and epoch == 0
                  and step < resume_step
              ):
                  continue
              outputs = model(**batch)
              loss = outputs.loss
              # We keep track of the loss at each epoch
              if args.with_tracking:
                  total_loss += loss.detach().float()
              loss = loss / args.gradient_accumulation_steps
              accelerator.backward(loss)
              if (
                  step % args.gradient_accumulation_steps == 0
                  or step == len(train_dataloader) - 1
              ):
                  optimizer.step()
                  lr_scheduler.step()
                  optimizer.zero_grad()
                  progress_bar.update(1)
                  completed_steps += 1
  
              if isinstance(checkpointing_steps, int):
                  if completed_steps % checkpointing_steps == 0:
                      output_dir = f"step_{completed_steps}"
                      if args.output_dir is not None:
                          output_dir = os.path.join(
                              args.output_dir, output_dir
                          )
                      accelerator.save_state(output_dir)
              if completed_steps >= args.max_train_steps:
                  break
  
          model.eval()
          losses = []
          for step, batch in enumerate(eval_dataloader):
              with torch.no_grad():
                  outputs = model(**batch)
  
              loss = outputs.loss
              losses.append(
                  accelerator.gather(
                      loss.repeat(args.per_device_eval_batch_size)
                  )
              )
  
          losses = torch.cat(losses)
          losses = losses[: len(eval_dataset)]
          try:
              perplexity = math.exp(torch.mean(losses))
          except OverflowError:
              perplexity = float("inf")
  
          logger.info(f"epoch {epoch}: perplexity: {perplexity}")
  
          if args.with_tracking:
              accelerator.log(
                  {
                      "perplexity": perplexity,
                      "train_loss": total_loss,
                      "epoch": epoch,
                      "step": completed_steps,
                  },
              )
  
          if args.push_to_hub and epoch < args.num_train_epochs - 1:
              accelerator.wait_for_everyone()
              unwrapped_model = accelerator.unwrap_model(model)
              unwrapped_model.save_pretrained(
                  args.output_dir, save_function=accelerator.save
              )
              if accelerator.is_main_process:
                  tokenizer.save_pretrained(args.output_dir)
                  repo.push_to_hub(
                      commit_message=f"Training in progress epoch {epoch}",
                      blocking=False,
                      auto_lfs_prune=True,
                  )
  
          if args.checkpointing_steps == "epoch":
              output_dir = f"epoch_{epoch}"
              if args.output_dir is not None:
                  output_dir = os.path.join(args.output_dir, output_dir)
              accelerator.save_state(output_dir)
  
      if args.output_dir is not None:
          accelerator.wait_for_everyone()
          unwrapped_model = accelerator.unwrap_model(model)
          unwrapped_model.save_pretrained(
              args.output_dir, save_function=accelerator.save
          )
          if accelerator.is_main_process:
              tokenizer.save_pretrained(args.output_dir)
              if args.push_to_hub:
                  repo.push_to_hub(
                      commit_message="End of training",
                      auto_lfs_prune=True,
                  )
  
          with open(
              os.path.join(args.output_dir, "all_results.json"), "w"
          ) as f:
              json.dump({"perplexity": perplexity}, f)
  
imports: import argparse
  import json
  import logging
  import math
  import os
  import random
  from itertools import chain
  from pathlib import Path
  import datasets
  import torch
  from datasets import load_dataset
  from torch.utils.data import DataLoader
  from tqdm.auto import tqdm
  import transformers
  from accelerate import Accelerator, DistributedType
  from accelerate.utils import set_seed
  from huggingface_hub import Repository
  from transformers import CONFIG_MAPPING, MODEL_MAPPING, AdamW, AutoConfig, AutoModelForCausalLM, AutoTokenizer, SchedulerType, default_data_collator, get_scheduler
  from transformers.utils import get_full_repo_name
  from transformers.utils.versions import require_version
label_desc: Use image node_id 2 for calling a global function with example usage: main() and returns: tokenizer, result