mkdataset.py

#! /usr/bin/env python

import argparse
import json

from mmlkg.data import dataset
from mmlkg.data.hdf5 import HDF5


_MODALITIES = ["textual", "numerical", "temporal", "visual", "spatial"]


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", help="Directory with CSV files "
                        + "(generated by `generateInput.py`)",
                        required=True)
    parser.add_argument("-c", "--config",
                        help="JSON file with hyperparameters",
                        default=None)
    parser.add_argument("-m", "--modalities", nargs='*',
                        help="Which modalities to include",
                        choices=[m.lower() for m in _MODALITIES],
                        default=_MODALITIES)
    parser.add_argument("-o", "--output", help="Output directory",
                        default=None)
    parser.add_argument("--compression", help="Compress generated dataset",
                        choices=["gzip", "lzf", "none"],
                        default="lzf")
    flags = parser.parse_args()

    config = {"encoders": dict(), "optim": dict()}
    if flags.config is not None:
        print("[CONF] Using configuration from %s" % flags.config)
        with open(flags.config, 'r') as f:
            config = json.load(f)

    out_dir = flags.input if flags.output is None else flags.output
    out_dir = out_dir + '/' if not out_dir.endswith('/') else out_dir

    out_file = out_dir + 'dataset.h5'
    with HDF5(out_file, 'w', flags.compression) as hf:
        nc_data = dict()
        lp_data = dict()
        for name, item in dataset.generate_dataset(flags, config):
            if name == 'num_nodes':
                hf.write_metadata(name, item)
            elif name in _MODALITIES:
                hf.write_modality_data(item, name)
            elif name in ['num_classes',
                          'training',
                          'testing',
                          'validation']:
                if item is not None:
                    nc_data[name] = item
            elif name in ['entities',
                          'triples',
                          'training_lp',
                          'testing_lp',
                          'validation_lp']:
                if item is not None:
                    lp_data[name] = item
            else:
                continue

        if len(nc_data) == 4:
            hf.write_task_data(nc_data, HDF5.NODE_CLASSIFICATION)
        if len(lp_data) == 5:
            hf.write_task_data(lp_data, HDF5.LINK_PREDICTION)

    print(f"Dataset saved to {out_file}")