Added weights to storage

marius-team · Dec 27, 2023 · c04d75a · c04d75a
1 parent ee171bb
commit c04d75a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/cpp/include/storage/graph_storage.h b/src/cpp/include/storage/graph_storage.h
@@ -13,8 +13,11 @@ struct GraphModelStoragePtrs {
     shared_ptr<Storage> edges = nullptr;
     shared_ptr<Storage> train_edges = nullptr;
     shared_ptr<Storage> train_edges_dst_sort = nullptr;
+    shared_ptr<Storage> train_edges_weights = nullptr;
     shared_ptr<Storage> validation_edges = nullptr;
+    shared_ptr<Storage> validation_edges_weights = nullptr;
     shared_ptr<Storage> test_edges = nullptr;
+    shared_ptr<Storage> test_edges_weights = nullptr;
     shared_ptr<Storage> nodes = nullptr;
     shared_ptr<Storage> train_nodes = nullptr;
     shared_ptr<Storage> valid_nodes = nullptr;

diff --git a/src/cpp/src/storage/io.cpp b/src/cpp/src/storage/io.cpp
@@ -486,6 +486,9 @@ shared_ptr<GraphModelStorage> initializeStorageLinkPrediction(shared_ptr<Model>
     storage_ptrs.train_edges_dst_sort = edge_storages["train_edges_dst_sort"];
     storage_ptrs.validation_edges = edge_storages["validation_edges"];
     storage_ptrs.test_edges = edge_storages["test_edges"];
+    storage_ptrs.train_edges_weights = edge_storages["train_edge_weights_storage"];
+    storage_ptrs.validation_edges_weights = edge_storages["valid_edge_weights_storage"];
+    storage_ptrs.test_edges_weights = edge_storages["test_edge_weights_storage"];
 
     storage_ptrs.node_features = initializeNodeFeatures(model, storage_config);
     storage_ptrs.node_embeddings = std::get<0>(node_embeddings);
@@ -506,10 +509,10 @@ shared_ptr<GraphModelStorage> initializeStorageNodeClassification(shared_ptr<Mod
     shared_ptr<Storage> node_labels = initializeNodeLabels(model, storage_config);
 
     GraphModelStoragePtrs storage_ptrs = {};
-
     storage_ptrs.train_edges = edge_storages["train_edges"];
     storage_ptrs.train_edges_dst_sort = edge_storages["train_edges_dst_sort"];
     storage_ptrs.edges = storage_ptrs.train_edges;
+    storage_ptrs.train_edges_weights = edge_storages["train_edge_weights_storage"];
 
     storage_ptrs.train_nodes = std::get<0>(node_id_storages);
     storage_ptrs.valid_nodes = std::get<1>(node_id_storages);