code: def get_model(training_args, config, tokenizer):
      # Find latest checkpoint in output_dir
      output_dir = Path(training_args.output_dir)
      logger.info(
          f'Checkpoint dir {output_dir}, contents {list(output_dir.glob("checkpoint*"))}'
      )
      latest_checkpoint_dir = max(
          output_dir.glob("checkpoint*"),
          default=None,
          key=os.path.getctime,
      )
  
      if latest_checkpoint_dir is not None:
          logger.info(f"Loading model from {latest_checkpoint_dir}")
          model = AlbertForPreTraining.from_pretrained(
              latest_checkpoint_dir
          )
      else:
          logger.info(f"Training from scratch")
          model = AlbertForPreTraining(config)
          model.resize_token_embeddings(len(tokenizer))
  
      return model
  
imports: import os
  import pickle
  import sys
  from dataclasses import asdict
  from pathlib import Path
  import torch
  import transformers
  from datasets import load_from_disk
  from torch.utils.data import DataLoader
  from torch_optimizer import Lamb
  from transformers import DataCollatorForLanguageModeling, HfArgumentParser, TrainingArguments, set_seed
  from transformers.models.albert import AlbertConfig, AlbertForPreTraining, AlbertTokenizerFast
  from transformers.optimization import get_linear_schedule_with_warmup
  from transformers.trainer import Trainer
  from transformers.trainer_utils import is_main_process
  from hivemind import DHT, Float16Compression, Optimizer, get_dht_time
  from hivemind.optim.state_averager import LRSchedulerBase
  from hivemind.utils.logging import get_logger, use_hivemind_log_handler
  from hivemind.utils.networking import log_visible_maddrs
  import utils
  from arguments import AlbertTrainingArguments, AveragerArguments, CollaborationArguments, DatasetArguments, ProgressTrackerArguments
label_desc: Use image node_id 2 for calling a global function with example usage: get_model(training_args, config, tokenizer) and returns: model