Source code for cogdl.wrappers.data_wrapper.link_prediction.gnn_kg_link_prediction_dw

from .. import DataWrapper


[docs]class GNNKGLinkPredictionDataWrapper(DataWrapper): def __init__(self, dataset): super(GNNKGLinkPredictionDataWrapper, self).__init__(dataset) self.dataset = dataset self.edge_set = None
[docs] def train_wrapper(self): return self.dataset.data
[docs] def val_wrapper(self): return self.dataset.data
[docs] def test_wrapper(self): return self.dataset.data