src/benchmarks/get_semistruct.py

import os.path as osp
from src.benchmarks.semistruct import AmazonSemiStruct, PrimeKGSemiStruct, MagSemiStruct


def get_semistructured_data(name, root='data/', download_processed=True, **kwargs):
    data_root = osp.join(root, name)
    if name == 'amazon':
        categories = ['Sports_and_Outdoors']
        kb = AmazonSemiStruct(root=data_root,
                                categories=categories,
                                meta_link_types=['brand'],
                                download_processed=download_processed,
                                **kwargs
                                )
    if name == 'primekg':
        kb = PrimeKGSemiStruct(root=data_root, 
                               download_processed=download_processed,
                               **kwargs)
    
    if name == 'mag':
        kb = MagSemiStruct(root=data_root, 
                           download_processed=download_processed)
    return kb