Wikidata5m knowledge graph dataset. Splits: train, valid, test
| 682 | |
| 683 | |
| 684 | class Wikidata5m(Dataset): |
| 685 | """ |
| 686 | Wikidata5m knowledge graph dataset. |
| 687 | |
| 688 | Splits: |
| 689 | train, valid, test |
| 690 | """ |
| 691 | def __init__(self): |
| 692 | super(Wikidata5m, self).__init__( |
| 693 | "wikidata5m", |
| 694 | urls={ |
| 695 | "train": "https://www.dropbox.com/s/dty6ufe1gg6keuc/wikidata5m.txt.gz?dl=1", |
| 696 | "valid": "https://www.dropbox.com/s/dty6ufe1gg6keuc/wikidata5m.txt.gz?dl=1", |
| 697 | "test": "https://www.dropbox.com/s/dty6ufe1gg6keuc/wikidata5m.txt.gz?dl=1", |
| 698 | "entity": "https://www.dropbox.com/s/bgmgvk8brjwpc9w/entity.txt.gz?dl=1", |
| 699 | "relation": "https://www.dropbox.com/s/37jxki93gguv0pp/relation.txt.gz?dl=1", |
| 700 | "alias2entity": [], # depends on `entity` |
| 701 | "alias2relation": [] # depends on `relation` |
| 702 | } |
| 703 | ) |
| 704 | |
| 705 | def train_preprocess(self, graph_file, train_file): |
| 706 | valid_file = train_file[:train_file.rfind("train.txt")] + "valid.txt" |
| 707 | test_file = train_file[:train_file.rfind("train.txt")] + "test.txt" |
| 708 | self.edge_split(graph_file, [train_file, valid_file, test_file], portions=[4000, 1, 1]) |
| 709 | |
| 710 | def valid_preprocess(self, graph_file, valid_file): |
| 711 | train_file = valid_file[:valid_file.rfind("valid.txt")] + "train.txt" |
| 712 | test_file = valid_file[:valid_file.rfind("valid.txt")] + "test.txt" |
| 713 | self.edge_split(graph_file, [train_file, valid_file, test_file], portions=[4000, 1, 1]) |
| 714 | |
| 715 | def test_preprocess(self, graph_file, test_file): |
| 716 | train_file = test_file[:test_file.rfind("valid.txt")] + "train.txt" |
| 717 | valid_file = test_file[:test_file.rfind("train.txt")] + "valid.txt" |
| 718 | self.edge_split(graph_file, [train_file, valid_file, test_file], portions=[4000, 1, 1]) |
| 719 | |
| 720 | def load_alias(self, alias_file): |
| 721 | alias2object = {} |
| 722 | ambiguous = set() |
| 723 | with open(alias_file, "r") as fin: |
| 724 | for line in fin: |
| 725 | tokens = line.strip().split("\t") |
| 726 | object = tokens[0] |
| 727 | for alias in tokens[1:]: |
| 728 | if alias in alias2object and alias2object[alias] != object: |
| 729 | ambiguous.add(alias) |
| 730 | alias2object[alias] = object |
| 731 | for alias in ambiguous: |
| 732 | alias2object.pop(alias) |
| 733 | return alias2object |
| 734 | |
| 735 | def alias2entity_preprocess(self, save_file): |
| 736 | return self.load_alias(self.entity) |
| 737 | |
| 738 | def alias2relation_preprocess(self, save_file): |
| 739 | return self.load_alias(self.relation) |
| 740 | |
| 741 |