code: def main(rank: int, world_size: int, args):
      """Args:
      ----
          rank (int): device id to put the model on
          world_size (int): number of processes participating in the job
          args: other arguments passed through the CLI
      """
      print(args)
  
      if rank == 0 and args.push_to_hub:
          login_to_hub()
  
      if not isinstance(args.workers, int):
          args.workers = min(16, multiprocessing.cpu_count())
  
      torch.backends.cudnn.benchmark = True
  
      vocab = VOCABS[args.vocab]
      fonts = args.font.split(",")
  
      if rank == 0:
          # Load val data generator
          st = time.time()
          if isinstance(args.val_path, str):
              with open(
                  os.path.join(args.val_path, "labels.json"), "rb"
              ) as f:
                  val_hash = hashlib.sha256(f.read()).hexdigest()
  
              val_set = RecognitionDataset(
                  img_folder=os.path.join(args.val_path, "images"),
                  labels_path=os.path.join(
                      args.val_path, "labels.json"
                  ),
                  img_transforms=T.Resize(
                      (args.input_size, 4 * args.input_size),
                      preserve_aspect_ratio=True,
                  ),
              )
          else:
              val_hash = None
              # Load synthetic data generator
              val_set = WordGenerator(
                  vocab=vocab,
                  min_chars=args.min_chars,
                  max_chars=args.max_chars,
                  num_samples=args.val_samples * len(vocab),
                  font_family=fonts,
                  img_transforms=Compose(
                      [
                          T.Resize(
                              (args.input_size, 4 * args.input_size),
                              preserve_aspect_ratio=True,
                          ),
                          # Ensure we have a 90% split of white-background images
                          T.RandomApply(T.ColorInversion(), 0.9),
                      ]
                  ),
              )
  
          val_loader = DataLoader(
              val_set,
              batch_size=args.batch_size,
              drop_last=False,
              num_workers=args.workers,
              sampler=SequentialSampler(val_set),
              pin_memory=torch.cuda.is_available(),
              collate_fn=val_set.collate_fn,
          )
          print(
              f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
              f"{len(val_loader)} batches)"
          )
  
      batch_transforms = Normalize(
          mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301)
      )
  
      # Load doctr model
      model = recognition.__dict__[args.arch](
          pretrained=args.pretrained, vocab=vocab
      )
  
      # Resume weights
      if isinstance(args.resume, str):
          print(f"Resuming {args.resume}")
          checkpoint = torch.load(args.resume, map_location="cpu")
          model.load_state_dict(checkpoint)
  
      # Backbone freezing
      if args.freeze_backbone:
          for p in model.feat_extractor.parameters():
              p.requires_grad = False
  
      # create default process group
      device = torch.device("cuda", args.devices[rank])
      dist.init_process_group(
          args.backend, rank=rank, world_size=world_size
      )
      # create local model
      model = model.to(device)
      # construct DDP model
      model = DDP(model, device_ids=[device])
  
      if rank == 0:
          # Metrics
          val_metric = TextMatch()
  
      if rank == 0 and args.test_only:
          print("Running evaluation")
          val_loss, exact_match, partial_match = evaluate(
              model,
              device,
              val_loader,
              batch_transforms,
              val_metric,
              amp=args.amp,
          )
          print(
              f"Validation loss: {val_loss:.6} (Exact: {exact_match:.2%} | Partial: {partial_match:.2%})"
          )
          return
  
      st = time.time()
  
      if isinstance(args.train_path, str):
          # Load train data generator
          base_path = Path(args.train_path)
          parts = (
              [base_path]
              if base_path.joinpath("labels.json").is_file()
              else [
                  base_path.joinpath(sub)
                  for sub in os.listdir(base_path)
              ]
          )
          with open(parts[0].joinpath("labels.json"), "rb") as f:
              train_hash = hashlib.sha256(f.read()).hexdigest()
  
          train_set = RecognitionDataset(
              parts[0].joinpath("images"),
              parts[0].joinpath("labels.json"),
              img_transforms=Compose(
                  [
                      T.Resize(
                          (args.input_size, 4 * args.input_size),
                          preserve_aspect_ratio=True,
                      ),
                      # Augmentations
                      T.RandomApply(T.ColorInversion(), 0.1),
                      RandomGrayscale(p=0.1),
                      RandomPhotometricDistort(p=0.1),
                      T.RandomApply(T.RandomShadow(), p=0.4),
                      T.RandomApply(
                          T.GaussianNoise(mean=0, std=0.1), 0.1
                      ),
                      T.RandomApply(GaussianBlur(3), 0.3),
                      RandomPerspective(distortion_scale=0.2, p=0.3),
                  ]
              ),
          )
          if len(parts) > 1:
              for subfolder in parts[1:]:
                  train_set.merge_dataset(
                      RecognitionDataset(
                          subfolder.joinpath("images"),
                          subfolder.joinpath("labels.json"),
                      )
                  )
      else:
          train_hash = None
          # Load synthetic data generator
          train_set = WordGenerator(
              vocab=vocab,
              min_chars=args.min_chars,
              max_chars=args.max_chars,
              num_samples=args.train_samples * len(vocab),
              font_family=fonts,
              img_transforms=Compose(
                  [
                      T.Resize(
                          (args.input_size, 4 * args.input_size),
                          preserve_aspect_ratio=True,
                      ),
                      # Ensure we have a 90% split of white-background images
                      T.RandomApply(T.ColorInversion(), 0.9),
                      RandomGrayscale(p=0.1),
                      RandomPhotometricDistort(p=0.1),
                      T.RandomApply(T.RandomShadow(), p=0.4),
                      T.RandomApply(
                          T.GaussianNoise(mean=0, std=0.1), 0.1
                      ),
                      T.RandomApply(GaussianBlur(3), 0.3),
                      RandomPerspective(distortion_scale=0.2, p=0.3),
                  ]
              ),
          )
  
      train_loader = DataLoader(
          train_set,
          batch_size=args.batch_size,
          drop_last=True,
          num_workers=args.workers,
          sampler=DistributedSampler(
              train_set,
              num_replicas=world_size,
              rank=rank,
              shuffle=True,
              drop_last=True,
          ),
          pin_memory=torch.cuda.is_available(),
          collate_fn=train_set.collate_fn,
      )
      print(
          f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
          f"{len(train_loader)} batches)"
      )
  
      if rank == 0 and args.show_samples:
          x, target = next(iter(train_loader))
          plot_samples(x, target)
          return
  
      # Optimizer
      optimizer = torch.optim.Adam(
          [p for p in model.parameters() if p.requires_grad],
          args.lr,
          betas=(0.95, 0.99),
          eps=1e-6,
          weight_decay=args.weight_decay,
      )
      # Scheduler
      if args.sched == "cosine":
          scheduler = CosineAnnealingLR(
              optimizer,
              args.epochs * len(train_loader),
              eta_min=args.lr / 25e4,
          )
      elif args.sched == "onecycle":
          scheduler = OneCycleLR(
              optimizer, args.lr, args.epochs * len(train_loader)
          )
  
      # Training monitoring
      current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      exp_name = (
          f"{args.arch}_{current_time}"
          if args.name is None
          else args.name
      )
  
      # W&B
      if rank == 0 and args.wb:
          run = wandb.init(
              name=exp_name,
              project="text-recognition",
              config={
                  "learning_rate": args.lr,
                  "epochs": args.epochs,
                  "weight_decay": args.weight_decay,
                  "batch_size": args.batch_size,
                  "architecture": args.arch,
                  "input_size": args.input_size,
                  "optimizer": "adam",
                  "framework": "pytorch",
                  "scheduler": args.sched,
                  "vocab": args.vocab,
                  "train_hash": train_hash,
                  "val_hash": val_hash,
                  "pretrained": args.pretrained,
              },
          )
  
      # Create loss queue
      min_loss = np.inf
      if args.early_stop:
          early_stopper = EarlyStopper(
              patience=args.early_stop_epochs,
              min_delta=args.early_stop_delta,
          )
      # Training loop
      for epoch in range(args.epochs):
          fit_one_epoch(
              model,
              device,
              train_loader,
              batch_transforms,
              optimizer,
              scheduler,
              amp=args.amp,
          )
  
          if rank == 0:
              # Validation loop at the end of each epoch
              val_loss, exact_match, partial_match = evaluate(
                  model,
                  device,
                  val_loader,
                  batch_transforms,
                  val_metric,
                  amp=args.amp,
              )
              if val_loss < min_loss:
                  # All processes should see same parameters as they all start from same
                  # random parameters and gradients are synchronized in backward passes.
                  # Therefore, saving it in one process is sufficient.
                  print(
                      f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state..."
                  )
                  torch.save(
                      model.module.state_dict(), f"./{exp_name}.pt"
                  )
              min_loss = val_loss
              print(
                  f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} "
                  f"(Exact: {exact_match:.2%} | Partial: {partial_match:.2%})"
              )
              # W&B
              if args.wb:
                  wandb.log(
                      {
                          "val_loss": val_loss,
                          "exact_match": exact_match,
                          "partial_match": partial_match,
                      }
                  )
              if args.early_stop and early_stopper.early_stop(val_loss):
                  print(
                      "Training halted early due to reaching patience limit."
                  )
                  break
      if rank == 0:
          if args.wb:
              run.finish()
  
          if args.push_to_hub:
              push_to_hf_hub(
                  model, exp_name, task="recognition", run_config=args
              )
  
imports: import os
  import datetime
  import hashlib
  import multiprocessing
  import time
  from pathlib import Path
  import numpy
  import torch
  import torch.distributed
  import torch.multiprocessing
  import wandb
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR
  from torch.utils.data import DataLoader, SequentialSampler
  from torch.utils.data.distributed import DistributedSampler
  from torchvision.transforms.v2 import Compose, GaussianBlur, Normalize, RandomGrayscale, RandomPerspective, RandomPhotometricDistort
  from tqdm.auto import tqdm
  from doctr import transforms
  from doctr.datasets import VOCABS, RecognitionDataset, WordGenerator
  from doctr.models import login_to_hub, push_to_hf_hub, recognition
  from doctr.utils.metrics import TextMatch
  from utils import EarlyStopper, plot_samples
label_desc: Use image node_id 3 for calling a global function with example usage: main(rank, world_size, args) without return types