code: def checkpoint_model(
      checkpoint_folder,
      ckpt_id,
      model,
      epoch,
      last_global_step,
      **kwargs,
  ):
      """Utility function for checkpointing model + optimizer dictionaries
      The main purpose for this is to be able to resume training from that instant again
      """
      checkpoint_state_dict = {
          "epoch": epoch,
          "last_global_step": last_global_step,
      }
      # Add extra kwargs too
      checkpoint_state_dict.update(kwargs)
  
      success = model.save_checkpoint(
          checkpoint_folder, ckpt_id, checkpoint_state_dict
      )
      status_msg = f"checkpointing: checkpoint_folder={checkpoint_folder}, ckpt_id={ckpt_id}"
      if success:
          logging.info(f"Success {status_msg}")
      else:
          logging.warning(f"Failure {status_msg}")
      return
  
imports: import argparse
  import json
  import logging
  import math
  import os
  import random
  from pathlib import Path
  from time import time
  import datasets
  import nltk
  import numpy
  import torch
  from datasets import load_dataset, load_metric
  from torch.utils.data import DataLoader
  from tqdm.auto import tqdm
  import transformers
  from accelerate import Accelerator
  from accelerate.logging import get_logger
  from accelerate.utils import DummyOptim, DummyScheduler, set_seed
  from filelock import FileLock
  from huggingface_hub import Repository
  from transformers import CONFIG_MAPPING, MODEL_MAPPING, AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, DataCollatorForSeq2Seq, SchedulerType, get_scheduler
  from transformers.utils import get_full_repo_name, is_offline_mode
  from transformers.utils.versions import require_version
label_desc: Use image node_id 2 for calling a global function with example usage: checkpoint_model(checkpoint_folder, ckpt_id, model, epoch, last_global_step) without return types