def get_dataset()

in python/dglke/dataloader/KGDataset.py [0:0]

• 31 lines of code
• 12 McCabe index (conditional complexity)

def get_dataset(data_path, data_name, format_str, delimiter='\t', files=None, has_edge_importance=False):
    if format_str == 'built_in':
        if data_name == 'Freebase':
            dataset = KGDatasetFreebase(data_path)
        elif data_name == 'FB15k':
            dataset = KGDatasetFB15k(data_path)
        elif data_name == 'FB15k-237':
            dataset = KGDatasetFB15k237(data_path)
        elif data_name == 'wn18':
            dataset = KGDatasetWN18(data_path)
        elif data_name == 'wn18rr':
            dataset = KGDatasetWN18rr(data_path)
        elif data_name == 'wikikg2':
            dataset = KGDatasetWikikg2(data_path)
        elif data_name == 'biokg':
            dataset = KGDatasetBiokg(data_path)
        elif data_name == 'wikikg90M':
            dataset = KGDatasetWiki90M(data_path)
        else:
            assert False, "Unknown dataset {}".format(data_name)
    elif format_str.startswith('raw_udd'):
        # user defined dataset
        assert data_name != 'FB15k', 'You should provide the dataset name for raw_udd format.'
        format = format_str[8:]
        dataset = KGDatasetUDDRaw(data_path, data_name, delimiter, files, format, has_edge_importance)
    elif format_str.startswith('udd'):
        # user defined dataset
        assert data_name != 'FB15k', 'You should provide the dataset name for udd format.'
        format = format_str[4:]
        dataset = KGDatasetUDD(data_path, data_name, delimiter, files, format, has_edge_importance)
    else:
        assert False, "Unknown format {}".format(format_str)

    return dataset