code: def get_input_data(stage, cpn_input_data):
      if stage == "train":
          train_data, validate_data = cpn_input_data
          train_data = _parse_data(train_data)
          if validate_data is not None:
              validate_data = _parse_data(validate_data)
          else:
              validate_data = None
  
          return train_data, validate_data
  
      elif stage == "predict":
          test_data = cpn_input_data
          test_data = _parse_data(test_data)
          return test_data
      else:
          raise ValueError(f"Unknown stage {stage}")
  
file_path: FATE/python/fate/components/components/nn/component_utils.py
imports: from fate.components.components.nn.loader import Loader
  from fate.components.components.nn.nn_runner import NNRunner
  from fate.components.components.utils import consts
  from fate.arch.dataframe import DataFrame
  from fate.components.components.utils.tools import add_dataset_type
  from fate.arch import Context
  from fate.components.core import ARBITER, GUEST, HOST, Role, cpn
  import logging
  from fate.components.core.component_desc.artifacts.data._dataframe import DataframeReader, DataframeWriter
  from fate.components.core.component_desc.artifacts.data._directory import DataDirectoryReader
  from fate.components.core.component_desc.artifacts.model._directory import ModelDirectoryReader, ModelDirectoryWriter
label_desc: Use image node_id 4 for calling a global function with example usage: get_input_data(stage, cpn_input_data) and returns: train_data,  validate_data, test_data