code: def train(
      model_name: str,
      train_data,
      val_data=None,
      test_data=None,
      config: dict = None,
      save: Union[bool, str, Path] = True,
      dry_run: bool = False,
  ) -> ModelWrapper:
      """
      Starts a training run and returns a wrapper around the model.
  
      Args:
          model_name (str): Name of the model to train
          train_data: Training data. Multiple formats available, see docs.
          val_data (optional): Validation data. Multiple formats available, see docs.
              Defaults to None.
          test_data (optional): Test data. Multiple formats available, see docs.
              Defaults to None.
          config (dict, optional): Configuration of hyperparameters. If None (default),
              some default hyperparameters for each model are used.
          save (Union[bool, str, Path], optional): Whether to store model and logs to disk
              (and to comet.ml if connected). If False, nothing will be saved. If True,
              artefacts will be saved to a timestamped directory in
              ./traintool-experiments. If a directory (str or Path) is given, artefacts
              will be saved there. Defaults to True.
          dry_run (bool, optional): If True, the model will only be trained for one batch
              to check that everything works. This will still save the model to
              disk; use save=False to prevent this. Defaults to False.
  
      Returns:
          ModelWrapper: A wrapper around the original model
      """
  
      # TODO: Disabled for now, so that we can pass parameters more dynamically (e.g. pass
      #   along sklearn model params without specifying them explicitly), think about if
      #   this is still required.
      # Get default config and fill up with values from config (checking that all keys
      # are correct)
      # default_config = model_wrapper_class.default_config(model_name)
      # if config is None:
      #     config = default_config
      # else:
      #     config = _update_config(default_config, config)
  
      start_time = datetime.now()
  
      with tempfile.TemporaryDirectory() as tmp_dir:  # only used when save=True
          if config is None:
              config = {}
  
          # Resolve model class (done here so no output dir is created if model_name is
          # invalid).
          model_wrapper_class = _resolve_model(model_name)
  
          # Create out_dir.
          experiment_id = (
              f"{start_time.strftime('%Y-%m-%d_%H-%M-%S')}_{model_name}"
          )
          if save is True:  # timestamped dir in ./traintool-experiments
              out_dir = project_dir / experiment_id
              out_dir.mkdir(parents=True, exist_ok=False)
          elif save is False:  # temporary dir
              # TODO: For convenience, we are just making out_dir a temporary directory
              #   here, which will be discarded. Instead, we shouldn't save anything at
              #   all because it might be more performant.
              out_dir = Path(tmp_dir)
          else:  # use save as dir
              out_dir = Path(save)
              out_dir.mkdir(parents=True, exist_ok=True)
  
          # Create yml file in out_dir with some information about the experiment.
          _write_info_file(
              out_dir,
              status="Running",
              model_name=model_name,
              config=config,
              start_time=str(start_time),
          )
  
          # Stream stdout to log file in out_dir.
          logger.add(
              out_dir / "stdout.log",
              format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {message}",
              backtrace=True,
              diagnose=True,
          )
  
          # Print some info.
          logger.info("  traintool experiment  ".center(80, "="))
          logger.info("ID:".ljust(12) + str(experiment_id))
          logger.info("Model:".ljust(12) + str(model_name))
          logger.info("Config:".ljust(12) + str(config))
          logger.info("Output dir:".ljust(12) + str(out_dir))
          if save is False:
              logger.info(
                  " " * 12
                  + "(temporary directory, will be automatically removed)"
              )
          if "api_key" in comet_config:
              if comet_config["project_name"] is None:
                  project = "Uncategorized Experiments"
              else:
                  project = comet_config["project_name"]
              logger.info(
                  "Comet.ml:".ljust(12)
                  + f"Enabled (project: {project})"
              )
          if save is not False:
              logger.info(
                  "Load via:".ljust(12)
                  + f'traintool.load("{experiment_id}")'
              )
          logger.info("=" * 80)
          if dry_run:
              logger.info(">>> THIS IS JUST A DRY RUN <<<")
              logger.info("")
  
          # Create tensorboard writer
          writer = _create_tensorboard_writer(out_dir=out_dir)
  
          # Create comet.ml experiment (or dummy object if comet is not used).
          # This has to be done right before training because it prints some stuff.
          experiment = _create_comet_experiment(
              config=config, save=save
          )
  
          try:
              # Create model wrapper and start training.
              model_wrapper = model_wrapper_class(
                  model_name, config, out_dir
              )
              model_wrapper._train(
                  train_data=train_data,
                  val_data=val_data,
                  test_data=test_data,
                  writer=writer,
                  experiment=experiment,
                  dry_run=dry_run,
              )
          except:  # noqa: E722
              status = "Failed"
              logger.exception(
                  "ERROR: Training interrupted by an exception (see below)"
              )
              raise
          else:
              status = "Finished"
          finally:
              # End experiment and write to log and info file.
              # TODO: Check if error is still tracked in comet ml experiment.
              experiment.end()
              writer.close()
              end_time = datetime.now()
              duration = end_time - start_time
              _write_info_file(
                  out_dir,
                  status=status,
                  end_time=str(end_time),
                  duration=str(duration),
              )
              logger.info(
                  f"  {status}! (after {utils.format_timedelta(duration)})  ".center(
                      80, "="
                  )
              )
  
          return model_wrapper
  
imports: from comet_ml import Experiment
  from typing import Union, Type
  from pathlib import Path
  import yaml
  from tensorboardX import SummaryWriter
  import tempfile
  import editdistance
  import numpy
  from datetime import datetime
  from loguru import logger
  import sys
  from .None import utils
  from .None import image_classification
  from .model_wrapper import ModelWrapper
label_desc: Use image node_id 7 for calling a global function with example usage: train(model_name, train_data, val_data, test_data, config, save, dry_run) and returns: model_wrapper