data.py 文件源码

def load(mode=DatasetMode.small):
    base_path = get_file(DataConstants.dataset, origin=DataConstants.origin, untar=True)
    base_path = os.path.join(base_path, mode)

    train_path = os.path.join(base_path, DataConstants.train)
    test_path = os.path.join(base_path, DataConstants.test)
    song_path = os.path.join(base_path, DataConstants.song_hash)

    songs = dict(read_song_hash(song_path))
    train, test = read_dataset(train_path, test_path)

    return train, test, songs