| 197 | return item |
| 198 | |
| 199 | def parse_graph_data(self): |
| 200 | |
| 201 | all_nodes = [] |
| 202 | all_edges = [] |
| 203 | all_edges_pad = [] |
| 204 | all_edges_ind = [] |
| 205 | |
| 206 | for g_ind, g in enumerate(self.graph): |
| 207 | G = nx.DiGraph() |
| 208 | g += ' ' |
| 209 | for p in g.split('__subject__')[1:]: |
| 210 | head = p.split('__predicate__')[0][1:-1] |
| 211 | relop = p.split('__predicate__')[1].split('__object__')[0][1:-1] |
| 212 | tail = p.split('__predicate__')[1].split('__object__')[1][1:-1] |
| 213 | G.add_edge(head, tail, edge=relop) |
| 214 | G.nodes[head]['node'] = head |
| 215 | G.nodes[tail]['node'] = tail |
| 216 | |
| 217 | nodes = list(G.nodes) + max(0, self.max_nodes-len(G.nodes)) * ['__no_node__'] |
| 218 | |
| 219 | edges = [] |
| 220 | edges_ind = [] |
| 221 | for u, v, d in G.edges(data=True): |
| 222 | edges.append(d['edge']) |
| 223 | edges_ind.append((nodes.index(u), nodes.index(v))) |
| 224 | |
| 225 | edges_pad = edges + max(0, self.max_edges - len(edges)) * ['__no_edge__'] |
| 226 | |
| 227 | all_nodes.append(nodes) |
| 228 | all_edges.append(edges) |
| 229 | all_edges_pad.append(edges_pad) |
| 230 | all_edges_ind.append(edges_ind) |
| 231 | |
| 232 | self.node = all_nodes |
| 233 | self.edge = all_edges |
| 234 | self.edge_pad = all_edges_pad |
| 235 | self.edge_ind = all_edges_ind |
| 236 | |
| 237 | def _build_inputs_with_special_tokens(self, token_ids_0, _): |
| 238 | # T5: <pad_id> token_ids_0 <eos_id> |