code: def load_env(data_dir, model_dir):
      """Loads environment for inference mode, used in jupyter notebook."""
      model_params = sketch_rnn_model.get_default_hparams()
      with tf.gfile.Open(
          os.path.join(model_dir, "model_config.json"), "r"
      ) as f:
          model_params.parse_json(f.read())
      return load_dataset(data_dir, model_params, inference_mode=True)
  
imports: import json
  import os
  import time
  import zipfile
  from magenta.models.sketch_rnn import model
  from magenta.models.sketch_rnn import utils
  import numpy
  import requests
  import six
  from six.moves.urllib.request import urlretrieve
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: load_env(data_dir, model_dir) and returns: load_dataset