(self)
| 283 | ) |
| 284 | |
| 285 | def process(self): |
| 286 | if self.split_ratio is None: |
| 287 | # Handle logics for OGB link prediction dataset |
| 288 | assert hasattr( |
| 289 | self.dataset, "get_edge_split" |
| 290 | ), "dataset doesn't have get_edge_split method, please specify split_ratio and neg_ratio to generate the split" |
| 291 | # This is likely to be an ogb dataset |
| 292 | self.edge_split = self.dataset.get_edge_split() |
| 293 | self._train_graph = self.g |
| 294 | if "source_node" in self.edge_split["test"]: |
| 295 | # Probably ogbl-citation2 |
| 296 | pos_e = ( |
| 297 | self.edge_split["valid"]["source_node"], |
| 298 | self.edge_split["valid"]["target_node"], |
| 299 | ) |
| 300 | neg_e_size = self.edge_split["valid"]["target_node_neg"].shape[ |
| 301 | -1 |
| 302 | ] |
| 303 | neg_e_src = np.repeat( |
| 304 | self.edge_split["valid"]["source_node"], neg_e_size |
| 305 | ) |
| 306 | neg_e_dst = np.reshape( |
| 307 | self.edge_split["valid"]["target_node_neg"], -1 |
| 308 | ) |
| 309 | self._val_edges = pos_e, (neg_e_src, neg_e_dst) |
| 310 | pos_e = ( |
| 311 | self.edge_split["test"]["source_node"], |
| 312 | self.edge_split["test"]["target_node"], |
| 313 | ) |
| 314 | neg_e_size = self.edge_split["test"]["target_node_neg"].shape[ |
| 315 | -1 |
| 316 | ] |
| 317 | neg_e_src = np.repeat( |
| 318 | self.edge_split["test"]["source_node"], neg_e_size |
| 319 | ) |
| 320 | neg_e_dst = np.reshape( |
| 321 | self.edge_split["test"]["target_node_neg"], -1 |
| 322 | ) |
| 323 | self._test_edges = pos_e, (neg_e_src, neg_e_dst) |
| 324 | elif "edge" in self.edge_split["test"]: |
| 325 | # Probably ogbl-collab |
| 326 | pos_e_tensor, neg_e_tensor = ( |
| 327 | self.edge_split["valid"]["edge"], |
| 328 | self.edge_split["valid"]["edge_neg"], |
| 329 | ) |
| 330 | pos_e = (pos_e_tensor[:, 0], pos_e_tensor[:, 1]) |
| 331 | neg_e = (neg_e_tensor[:, 0], neg_e_tensor[:, 1]) |
| 332 | self._val_edges = pos_e, neg_e |
| 333 | |
| 334 | pos_e_tensor, neg_e_tensor = ( |
| 335 | self.edge_split["test"]["edge"], |
| 336 | self.edge_split["test"]["edge_neg"], |
| 337 | ) |
| 338 | pos_e = (pos_e_tensor[:, 0], pos_e_tensor[:, 1]) |
| 339 | neg_e = (neg_e_tensor[:, 0], neg_e_tensor[:, 1]) |
| 340 | self._test_edges = pos_e, neg_e |
| 341 | # delete edge split to save memory |
| 342 | self.edge_split = None |
nothing calls this directly
no test coverage detected