| 74 | |
| 75 | |
| 76 | class CoraDataset(Dataset): |
| 77 | CONTENT_FILE_NAME = "cora.content" |
| 78 | CITE_FILE_NAME = "cora.cites" |
| 79 | |
| 80 | def __init__(self, input_path, train_ratio): |
| 81 | Dataset.__init__(self, input_path) |
| 82 | self.train_ratio = train_ratio |
| 83 | self.content_file = os.path.join(self.input_path, |
| 84 | CoraDataset.CONTENT_FILE_NAME) |
| 85 | self.cite_file = os.path.join(self.input_path, |
| 86 | CoraDataset.CITE_FILE_NAME) |
| 87 | self.node_label_feature = [] |
| 88 | self.edges = [] |
| 89 | self.node_id_map = {} |
| 90 | |
| 91 | def load(self): |
| 92 | self.node_label_feature = self.__load_node_label_feature() |
| 93 | self.edges = self.__load_edges() |
| 94 | |
| 95 | def to_nx_graph(self): |
| 96 | self.node_id_map = self.__build_node_id_map() |
| 97 | |
| 98 | train_num = len(self.node_id_map) * self.train_ratio |
| 99 | |
| 100 | nx_graph = nx.Graph() |
| 101 | for node, label, feature in self.node_label_feature: |
| 102 | node_id = self.node_id_map[node] |
| 103 | stage = "train" if node_id <= train_num else "test" |
| 104 | feature = np.array(feature) / (np.sum(feature) + 1e-7) |
| 105 | nx_graph.add_node(node_id, |
| 106 | stage=stage, |
| 107 | feature=feature, |
| 108 | label=label) |
| 109 | |
| 110 | for src_node, dst_node in self.edges: |
| 111 | if src_node in self.node_id_map and dst_node in self.node_id_map: |
| 112 | src_id = self.node_id_map[src_node] |
| 113 | dst_id = self.node_id_map[dst_node] |
| 114 | if src_id <= train_num and dst_id <= train_num: |
| 115 | stage = "train" |
| 116 | else: |
| 117 | stage = "test" |
| 118 | nx_graph.add_edge(src_id, dst_id, stage=stage, weight=1.0) |
| 119 | return nx_graph |
| 120 | |
| 121 | def __load_node_label_feature(self): |
| 122 | node_label_feature = [] |
| 123 | label_id_map = {} |
| 124 | with open(self.content_file, 'r', encoding='utf-8') as fin: |
| 125 | for line in fin: |
| 126 | fields = line.rstrip().split("\t") |
| 127 | node, label = fields[0], fields[-1] |
| 128 | feature = [int(x) for x in fields[1:-1]] |
| 129 | if label not in label_id_map: |
| 130 | label_id_map[label] = len(label_id_map) |
| 131 | node_label_feature.append((node, label_id_map[label], feature)) |
| 132 | return node_label_feature |
| 133 | |