code: def main(args, model=None) -> SummarizationModule:
      Path(args.output_dir).mkdir(exist_ok=True)
      check_output_dir(args, expected_items=3)
  
      if model is None:
          if "summarization" in args.task:
              model: SummarizationModule = SummarizationModule(args)
          else:
              model: SummarizationModule = TranslationModule(args)
      dataset = Path(args.data_dir).name
      if (
          args.logger_name == "default"
          or args.fast_dev_run
          or str(args.output_dir).startswith("/tmp")
          or str(args.output_dir).startswith("/var")
      ):
          logger = True  # don't pollute wandb logs unnecessarily
      elif args.logger_name == "wandb":
          from pytorch_lightning.loggers import WandbLogger
  
          project = os.environ.get("WANDB_PROJECT", dataset)
          logger = WandbLogger(
              name=model.output_dir.name, project=project
          )
  
      elif args.logger_name == "wandb_shared":
          from pytorch_lightning.loggers import WandbLogger
  
          logger = WandbLogger(
              name=model.output_dir.name, project=f"hf_{dataset}"
          )
  
      if args.early_stopping_patience >= 0:
          es_callback = get_early_stopping_callback(
              model.val_metric, args.early_stopping_patience
          )
      else:
          es_callback = False
  
      lower_is_better = args.val_metric == "loss"
      trainer: pl.Trainer = generic_train(
          model,
          args,
          logging_callback=Seq2SeqLoggingCallback(),
          checkpoint_callback=get_checkpoint_callback(
              args.output_dir,
              model.val_metric,
              args.save_top_k,
              lower_is_better,
          ),
          early_stopping_callback=es_callback,
          logger=logger,
      )
      pickle_save(model.hparams, model.output_dir / "hparams.pkl")
      if not args.do_predict:
          return model
  
      model.hparams.test_checkpoint = ""
      checkpoints = list(
          sorted(
              glob.glob(
                  os.path.join(args.output_dir, "*.ckpt"),
                  recursive=True,
              )
          )
      )
      if checkpoints:
          model.hparams.test_checkpoint = checkpoints[-1]
          trainer.resume_from_checkpoint = checkpoints[-1]
      trainer.logger.log_hyperparams(model.hparams)
  
      # test() without a model tests using the best checkpoint automatically
      trainer.test()
      return model
  
imports: import argparse
  import glob
  import logging
  import os
  import sys
  import time
  from collections import defaultdict
  from pathlib import Path
  from typing import Dict, List, Tuple
  import numpy
  import pytorch_lightning
  import torch
  from torch import nn
  from torch.utils.data import DataLoader
  from callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
  from transformers import MBartTokenizer, T5ForConditionalGeneration
  from transformers.models.bart.modeling_bart import shift_tokens_right
  from utils import ROUGE_KEYS, LegacySeq2SeqDataset, Seq2SeqDataset, assert_all_frozen, calculate_bleu, calculate_rouge, check_output_dir, flatten_list, freeze_embeds, freeze_params, get_git_info, label_smoothed_nll_loss, lmap, pickle_save, save_git_info, save_json, use_task_specific_params
  from lightning_base import BaseTransformer, add_generic_args, generic_train
label_desc: Use image node_id 1 for calling a global function with example usage: main(args, model) and returns: model, model