| 456 | |
| 457 | |
| 458 | class CitationDataset(InMemoryDataset): |
| 459 | ''' |
| 460 | ''' |
| 461 | file_names = {"Cora": "Cora.pt", |
| 462 | "Pubmed": "Pubmed.pt", "Citeseer": "Citeseer.pt"} |
| 463 | |
| 464 | def __init__(self, root, name, args, transform=None, pre_transform=None): |
| 465 | |
| 466 | self.name = name |
| 467 | assert self.name in ['Cora', 'Pubmed', 'Citeseer'] |
| 468 | |
| 469 | self.args = args |
| 470 | super().__init__(root, transform, pre_transform) |
| 471 | # if self.name == 'Citeseer': |
| 472 | # self.data = torch.load(osp.join(self.processed_dir, 'data.pt')) |
| 473 | # else: |
| 474 | |
| 475 | self.load(self.processed_paths[0]) |
| 476 | |
| 477 | if args.if_text: |
| 478 | if not osp.exists(osp.join(self.processed_dir, 'data_text.pt')): |
| 479 | self.text_emb() |
| 480 | torch.save(self.data, osp.join( |
| 481 | self.processed_dir, 'data_text.pt')) |
| 482 | else: |
| 483 | self.data = torch.load( |
| 484 | osp.join(self.processed_dir, 'data_text.pt')) |
| 485 | if self.name == 'Citeseer': |
| 486 | self.data.edge_index = torch.load( |
| 487 | "datasets/Citeseer/raw/data.pt").edge_index |
| 488 | self.label_emb_func(root) |
| 489 | |
| 490 | def label_emb_func(self, root): |
| 491 | if self.name == "Cora": |
| 492 | category_desc = pd.read_csv( |
| 493 | os.path.join(root + "/Cora", "categories.csv"), sep="," |
| 494 | ).values |
| 495 | ordered_desc = [] |
| 496 | label_names = self.data.label_names |
| 497 | for i, label in enumerate(label_names): |
| 498 | true_ind = label == category_desc[:, 0] |
| 499 | ordered_desc.append((label, category_desc[:, 1][true_ind])) |
| 500 | self.data.label_text = [ |
| 501 | desc[0] |
| 502 | + "." |
| 503 | + desc[1][0] |
| 504 | for desc in ordered_desc |
| 505 | ] |
| 506 | elif self.name == "Pubmed": |
| 507 | with open( |
| 508 | os.path.join(root + "/Pubmed", "categories.csv"), "r" |
| 509 | ) as f: |
| 510 | ordered_desc = f.read().split("\n") |
| 511 | self.data.label_text = [ |
| 512 | desc |
| 513 | for desc in ordered_desc |
| 514 | ] |
| 515 | elif self.name == "Citeseer": |
nothing calls this directly
no outgoing calls
no test coverage detected