Create a DGL Homogeneous graph with heterograph info stored as node or edge features.
(
num_nodes, num_rels, train, valid, test, reverse=True
)
| 274 | |
| 275 | |
| 276 | def build_knowledge_graph( |
| 277 | num_nodes, num_rels, train, valid, test, reverse=True |
| 278 | ): |
| 279 | """Create a DGL Homogeneous graph with heterograph info stored as node or edge features.""" |
| 280 | src = [] |
| 281 | rel = [] |
| 282 | dst = [] |
| 283 | raw_subg = {} |
| 284 | raw_subg_eset = {} |
| 285 | raw_subg_etype = {} |
| 286 | raw_reverse_sugb = {} |
| 287 | raw_reverse_subg_eset = {} |
| 288 | raw_reverse_subg_etype = {} |
| 289 | |
| 290 | # here there is noly one node type |
| 291 | s_type = "node" |
| 292 | d_type = "node" |
| 293 | |
| 294 | def add_edge(s, r, d, reverse, edge_set): |
| 295 | r_type = str(r) |
| 296 | e_type = (s_type, r_type, d_type) |
| 297 | if raw_subg.get(e_type, None) is None: |
| 298 | raw_subg[e_type] = ([], []) |
| 299 | raw_subg_eset[e_type] = [] |
| 300 | raw_subg_etype[e_type] = [] |
| 301 | raw_subg[e_type][0].append(s) |
| 302 | raw_subg[e_type][1].append(d) |
| 303 | raw_subg_eset[e_type].append(edge_set) |
| 304 | raw_subg_etype[e_type].append(r) |
| 305 | |
| 306 | if reverse is True: |
| 307 | r_type = str(r + num_rels) |
| 308 | re_type = (d_type, r_type, s_type) |
| 309 | if raw_reverse_sugb.get(re_type, None) is None: |
| 310 | raw_reverse_sugb[re_type] = ([], []) |
| 311 | raw_reverse_subg_etype[re_type] = [] |
| 312 | raw_reverse_subg_eset[re_type] = [] |
| 313 | raw_reverse_sugb[re_type][0].append(d) |
| 314 | raw_reverse_sugb[re_type][1].append(s) |
| 315 | raw_reverse_subg_eset[re_type].append(edge_set) |
| 316 | raw_reverse_subg_etype[re_type].append(r + num_rels) |
| 317 | |
| 318 | for edge in train: |
| 319 | s, r, d = edge |
| 320 | assert r < num_rels |
| 321 | add_edge(s, r, d, reverse, 1) # train set |
| 322 | |
| 323 | for edge in valid: |
| 324 | s, r, d = edge |
| 325 | assert r < num_rels |
| 326 | add_edge(s, r, d, reverse, 2) # valid set |
| 327 | |
| 328 | for edge in test: |
| 329 | s, r, d = edge |
| 330 | assert r < num_rels |
| 331 | add_edge(s, r, d, reverse, 3) # test set |
| 332 | |
| 333 | subg = [] |