code: def run(
      config_map,
      tf_file_reader=tf.data.TFRecordDataset,
      file_reader=tf.python_io.tf_record_iterator,
  ):
      """Load model params, save config file and start trainer.
  
      Args:
        config_map: Dictionary mapping configuration name to Config object.
        tf_file_reader: The tf.data.Dataset class to use for reading files.
        file_reader: The Python reader to use for reading files.
  
      Raises:
        ValueError: if required flags are missing or invalid.
      """
      if not FLAGS.run_dir:
          raise ValueError("Invalid run directory: %s" % FLAGS.run_dir)
      run_dir = os.path.expanduser(FLAGS.run_dir)
      train_dir = os.path.join(run_dir, "train")
  
      if FLAGS.mode not in ["train", "eval"]:
          raise ValueError("Invalid mode: %s" % FLAGS.mode)
  
      if FLAGS.config not in config_map:
          raise ValueError("Invalid config: %s" % FLAGS.config)
      config = config_map[FLAGS.config]
      if FLAGS.hparams:
          config.hparams.parse(FLAGS.hparams)
      config_update_map = {}
      if FLAGS.examples_path:
          config_update_map[
              "%s_examples_path" % FLAGS.mode
          ] = os.path.expanduser(FLAGS.examples_path)
      if FLAGS.tfds_name:
          if FLAGS.examples_path:
              raise ValueError(
                  "At most one of --examples_path and --tfds_name can be set."
              )
          config_update_map["tfds_name"] = FLAGS.tfds_name
          config_update_map["eval_examples_path"] = None
          config_update_map["train_examples_path"] = None
      config = configs.update_config(config, config_update_map)
      if FLAGS.num_sync_workers:
          config.hparams.batch_size //= FLAGS.num_sync_workers
  
      if FLAGS.mode == "train":
          is_training = True
      elif FLAGS.mode == "eval":
          is_training = False
      else:
          raise ValueError("Invalid mode: {}".format(FLAGS.mode))
  
      def dataset_fn():
          return data.get_dataset(
              config,
              tf_file_reader=tf_file_reader,
              is_training=is_training,
              cache_dataset=FLAGS.cache_dataset,
          )
  
      if is_training:
          train(
              train_dir,
              config=config,
              dataset_fn=dataset_fn,
              checkpoints_to_keep=FLAGS.checkpoints_to_keep,
              keep_checkpoint_every_n_hours=FLAGS.keep_checkpoint_every_n_hours,
              num_steps=FLAGS.num_steps,
              master=FLAGS.master,
              num_sync_workers=FLAGS.num_sync_workers,
              num_ps_tasks=FLAGS.num_ps_tasks,
              task=FLAGS.task,
          )
      else:
          num_batches = (
              FLAGS.eval_num_batches
              or data.count_examples(
                  config.eval_examples_path,
                  config.tfds_name,
                  config.data_converter,
                  file_reader,
              )
              // config.hparams.batch_size
          )
          eval_dir = os.path.join(
              run_dir, "eval" + FLAGS.eval_dir_suffix
          )
          evaluate(
              train_dir,
              eval_dir,
              config=config,
              dataset_fn=dataset_fn,
              num_batches=num_batches,
              master=FLAGS.master,
          )
  
file_path: magenta/magenta/models/music_vae/music_vae_train.py
imports: import os
  from magenta.models.music_vae import configs
  from magenta.models.music_vae import data
  import tensorflow.compat.v1
  import tf_slim
label_desc: Use image node_id 5 for calling a global function with example usage: run(config_map, tf_file_reader, file_reader) and returns: data