graph_sampler.py (14356B)
1 from rdflib.extras.external_graph_libs import rdflib_to_networkx_multidigraph 2 from rdflib.term import Literal 3 from typing import Any 4 import argparse 5 import collections 6 import json 7 import networkx as nx 8 import numpy 9 import pathlib 10 import random 11 import rdflib 12 import re 13 import sys 14 import tqdm 15 16 17 ENTITY_NAME_REGEX = re.compile(r'([A-Za-z]+)[0-9]+') 18 19 20 def get_entity_type(entity_name: str) -> str: 21 """ 22 Get the type of an entity. 23 24 This transforms "UndergraduateStudent143" into "UndergraduateStudent". 25 """ 26 return ENTITY_NAME_REGEX.fullmatch(entity_name)[1] 27 28 29 def parse_args() -> argparse.Namespace: 30 """ Return namespace containing command line arguments. """ 31 parser = argparse.ArgumentParser(description="Compositional reasoning over KGs and text using LM+GNNs: graph sampler.") 32 parser.add_argument("--graph-output-path", type=pathlib.Path, default=None, help="Graphs output folder path.") 33 parser.add_argument("--question-output-path", type=pathlib.Path, default=pathlib.Path("questions"), help="Questions output folder path.") 34 parser.add_argument("--input-path", type=pathlib.Path, default=pathlib.Path("raw_graphs"), help="Location of the input graphs generated by LUBM.") 35 parser.add_argument("--output-file-name", type=str, default="qs", help="Question file name.") 36 parser.add_argument("--entities-distribution", type=str, default="uniform", help="How to sample points for positive and negative (one of: vertex, degree, vertex-then-degree, walks).") 37 parser.add_argument("--entities-selection", type=str, default="unconditional", help="How to select the tail entity for positive and negative. Either unconditional or path.") 38 parser.add_argument("-k", "--hops", type=int, default=3, help="The number of hops in the generated paths.") 39 parser.add_argument("-n", "--samples", type=int, default=3, help="The number of samples to generate.") 40 return parser.parse_args() 41 42 43 def get_entities_by_type(KG: nx.MultiDiGraph) -> dict[str, list[str]]: 44 """ Get all vertices in a graph grouped by type. """ 45 entities_by_type: dict[str, set[str]] = collections.defaultdict(set) 46 for entity in KG.nodes(): 47 entities_by_type[get_entity_type(entity)].add(entity) 48 return {type: list(entities) for type, entities in entities_by_type.items()} 49 50 51 def get_degree_weighted_entities_by_type(KG: nx.MultiDiGraph) -> dict[str, list[str]]: 52 """ For each entity type, get a list of vertices of this type where each vertex appears as many time as its degree. """ 53 entities_by_type: dict[str, list[str]] = collections.defaultdict(list) 54 for edge in KG.edges(): 55 entities_by_type[get_entity_type(edge[0])].append(edge[0]) 56 return entities_by_type 57 58 59 def preprocess_graph(source: nx.MultiDiGraph) -> tuple[nx.MultiDiGraph, set[str]]: 60 """ 61 Symmetrize graph, filter irrelevant vertices and relations. 62 63 Literal entities such as phone numbers and RDF schema metadata are filtered out. 64 All vertex and relation labels are converted to standard strings. 65 For all arcs u->v labeled REL, a new arc v->u is created with label REL_r. 66 67 Returns: 68 nx.MultiDiGraph: the filtered graph 69 set[str]: the set of relations appearing in the graph 70 """ 71 relations: set[str] = set() 72 KG = nx.MultiDiGraph() 73 for i, (rdf_head, rdf_tail, rdf_pred) in enumerate(source.edges): 74 if isinstance(rdf_head, Literal) or isinstance(rdf_tail, Literal): 75 continue 76 head: str = str(rdf_head).split('/')[-1] 77 tail: str = str(rdf_tail).split('/')[-1] 78 pred: str = str(rdf_pred).split('#')[-1] 79 if 'type' in pred or 'import' in pred: 80 continue 81 if 'www' in head or 'www' in tail: 82 continue 83 pred_r: str = pred + '_r' 84 KG.add_edge(head, tail, type=pred, key=pred) 85 KG.add_edge(tail, head, type=pred_r, key=pred_r) 86 relations.add(pred_r) 87 relations.add(pred) 88 return KG, relations 89 90 91 def get_spelled_out_path(KG: Any, list_of_vertices: list[str], only_relations: bool) -> list[Any]: 92 """ Transform a list of vertices into a list of (head, relation, tail) or a list relations. """ 93 path: list[str | list[str]] = [] 94 for i in range(len(list_of_vertices)-1): 95 head: str = list_of_vertices[i] 96 tail: str = list_of_vertices[i+1] 97 kg_edge: dict[str, Any] = KG.get_edge_data(head, tail) 98 predicate: str 99 if isinstance(KG, nx.MultiDiGraph): 100 # In our "pruned" version of LUBM, there are no parallel edges going in the same direction. 101 predicate = list(kg_edge.keys())[0] 102 else: # nx.DiGraph 103 predicate = kg_edge["type"] 104 path.append(predicate if only_relations else [head, predicate, tail]) 105 return path 106 107 108 def select_probabilities(weights: dict[str, float], vertices: list[str]) -> numpy.ndarray: 109 """ Given the weight of all vertices, construct the probability vector associated with a given subset of vertices. """ 110 output: numpy.ndarray = numpy.empty(len(vertices), float) 111 for i, vertex in enumerate(vertices): 112 output[i] = weights[vertex] 113 return output / output.sum() 114 115 116 def sample_weighted_entities(KG: nx.MultiDiGraph, entities_by_type: dict[str, list[str]], hops: int, weights: list[dict[str, float]]) -> tuple[tuple[str, str], tuple[str, str]]: 117 """ 118 Sample 4 vertices from a categorical distribution, making sure they can work as positives and negatives. 119 120 Returns: (start1, end1), (start2, end2) 121 The points satisfy the following properties: 122 - start1 ≠ start2 or end1 ≠ end2 123 - type(start1) = type(start2) and type(end1) = type(end2) 124 """ 125 vertices: list[str] = list(KG.nodes()) 126 start1: str 127 end1: str 128 start1, end1 = numpy.random.choice(vertices, size=2, replace=False, p=select_probabilities(weights[0], vertices)) 129 start_type: str = get_entity_type(start1) 130 end_type: str = get_entity_type(end1) 131 132 start2: str = numpy.random.choice(entities_by_type[start_type], p=select_probabilities(weights[0], entities_by_type[start_type])) 133 end2: str 134 if start1 == start2: 135 # Make sure end1 ≠ end2 136 end2_weights: numpy.ndarray = numpy.array([0 if vertex == end1 else weights[0][vertex] for vertex in entities_by_type[end_type]], dtype=float) 137 end2_weights /= end2_weights.sum() 138 end2 = numpy.random.choice(entities_by_type[end_type], p=end2_weights) 139 else: 140 end2 = numpy.random.choice(entities_by_type[end_type], p=select_probabilities(weights[0], entities_by_type[end_type])) 141 142 return (start1, end1), (start2, end2) 143 144 145 def sample_paths_endpoints(KG: nx.MultiDiGraph, entities_by_type: dict[str, list[str]], hops: int, weights: list[dict[str, float]]) -> tuple[tuple[str, str], tuple[str, str]]: 146 """ Sample two paths and returns the endpoints, making sure these endpoints can work as positives and negatives. """ 147 def sample_path_endpoints(start_set: list[str]) -> tuple[None | str, None | str]: 148 path: list[str] = [numpy.random.choice(start_set, p=select_probabilities(weights[0], start_set))] 149 150 for hop in range(hops): 151 previous_vertex: str = path[-1] 152 candidates: list[str] = [vertex for vertex in KG.successors(previous_vertex) if vertex not in path] 153 if not candidates: 154 return None, None 155 next_vertex: str = numpy.random.choice(candidates, p=select_probabilities(weights[hop+1], candidates)) 156 path.append(next_vertex) 157 158 return path[0], path[-1] 159 160 vertices: list[str] = list(KG.nodes()) 161 while True: 162 start1, end1 = sample_path_endpoints(vertices) 163 if start1 is not None: 164 start2, end2 = sample_path_endpoints(entities_by_type[get_entity_type(start1)]) 165 if start2 is not None and start1 != start2 or start2 != end2: 166 return (start1, end1), (start2, end2) 167 168 169 def compute_uniform_weights(KG: nx.MultiDiGraph, hops: int) -> list[dict[str, float]]: 170 """ Make all vertices equally likely. """ 171 return [{vertex: 1 for vertex in KG.nodes()}] * (hops + 1) 172 173 174 def compute_degree_weights(KG: nx.MultiDiGraph, hops: int) -> list[dict[str, float]]: 175 """ Select vertices according to their degree. """ 176 return [{vertex: KG.degree(vertex) for vertex in KG.nodes()}] * (hops + 1) 177 178 179 def compute_degree_then_uniform_weights(KG: nx.MultiDiGraph, hops: int) -> list[dict[str, float]]: 180 """ Start selecting vertices according to their degree, then select neighbors uniformly. """ 181 return [{vertex: KG.degree(vertex) for vertex in KG.nodes()}] + [{vertex: 1 for vertex in KG.nodes()}] * hops 182 183 184 def compute_walks_weights(KG: nx.MultiDiGraph, hops: int) -> list[dict[str, float]]: 185 """ Select vertices according to the number of walks starting there. """ 186 vertices: list[str] = list(KG.nodes()) 187 adj: numpy.ndarray = nx.to_numpy_array(KG, nodelist=vertices) 188 walks: list[numpy.ndarray] = [numpy.ones(adj.shape[0])] 189 for hop in range(hops): 190 walks.append(walks[-1] @ adj) 191 return [dict(zip(vertices, x)) for x in walks[::-1]] 192 193 194 def entities_sampler(args: argparse.Namespace, KG: nx.MultiDiGraph, entities_by_type: dict[str, list[str]], hops: int) -> tuple[tuple[str, str], tuple[str, str]]: 195 """ Sample positive and negative entities. """ 196 weights: list[dict[str, float]] 197 if args.entities_distribution == "uniform": 198 weights = compute_uniform_weights(KG, hops) 199 elif args.entities_distribution == "degree": 200 weights = compute_degree_weights(KG, hops) 201 elif args.entities_distribution == "degree-then-uniform": 202 weights = compute_degree_then_uniform_weights(KG, hops) 203 elif args.entities_distribution == "walks": 204 weights = compute_walks_weights(KG, hops) 205 else: 206 print(f"Unknown entities-distribution: {args.entities_distribution}", file=sys.stderr) 207 sys.exit(1) 208 209 if args.entities_selection == "unconditional": 210 return sample_weighted_entities(KG, entities_by_type, hops, weights) 211 elif args.entities_selection == "path": 212 return sample_paths_endpoints(KG, entities_by_type, hops, weights) 213 else: 214 print(f"Unknown entities-selection: {args.entities_selection}", file=sys.stderr) 215 sys.exit(1) 216 217 218 def all_path_types(KG: nx.MultiDiGraph, start: str, end: str, hops: int) -> set[tuple[str, ...]]: 219 """ Get the set of all path types between two given vertices. """ 220 alternative_paths: list[list[str]] = [n for n in nx.all_simple_paths(KG, start, end, hops+1) if len(n) == hops+1] 221 rel_types_path: set[tuple[str, ...]] = set() 222 for alternative in alternative_paths: 223 relation_path: list[str] = get_spelled_out_path(KG, alternative, only_relations=True) 224 rel_types_path.add(tuple(relation_path)) 225 return rel_types_path 226 227 228 def sample_question(args: argparse.Namespace, KG: nx.MultiDiGraph, entities_by_type: dict[str, list[str]]) -> None | tuple[tuple[str, ...], tuple[str, str], tuple[str, str]]: 229 """ Sample a question with positive and negative pairs of entities. """ 230 while True: 231 (start1, end1), (start2, end2) = entities_sampler(args, KG, entities_by_type, args.hops) 232 path_types1: set[tuple[str, ...]] = all_path_types(KG, start1, end1, args.hops) 233 path_types2: set[tuple[str, ...]] = all_path_types(KG, start2, end2, args.hops) 234 235 type1_exclusives: set[tuple[str, ...]] = path_types1.difference(path_types2) 236 type2_exclusives: set[tuple[str, ...]] = path_types2.difference(path_types1) 237 # Make sure positives and negatives follow the same distribution 238 if type1_exclusives and type2_exclusives: 239 question: tuple[str, ...] = random.choice(list(type1_exclusives)) 240 return question, (start1, end1), (start2, end2) 241 242 243 def create_permutation(KG: nx.MultiDiGraph, entities_by_type: dict[str, list[str]]) -> dict[str, str]: 244 """ Create a type-preserving permutation of vertex labels. """ 245 mapping: dict[str, str] = {} 246 for entity_type, possibilities in entities_by_type.items(): 247 options: list[str] = possibilities 248 random.shuffle(options) 249 mapping.update(zip(possibilities, options)) 250 return mapping 251 252 253 if __name__ == '__main__': 254 args = parse_args() 255 256 # Create output directories if they do not exist 257 for output_directory in [args.graph_output_path, args.question_output_path]: 258 if not output_directory.is_dir(): 259 output_directory.mkdir(parents=True, exist_ok=True) 260 261 question_path: pathlib.Path = args.question_output_path / f'{args.hops}_hop_{args.output_file_name}_{args.samples}.jsonl' 262 with question_path.open("w") as question_file: 263 for _, input_file in zip(tqdm.trange(args.samples//2), args.input_path.iterdir()): 264 graph_path: pathlib.Path = args.graph_output_path / f'{input_file.stem}.graphml' 265 266 rdf_graph = rdflib.Graph() 267 rdf_graph.parse(input_file) 268 original_graph: nx.MultiDiGraph = rdflib_to_networkx_multidigraph(rdf_graph) 269 filtered_graph, possible_relations = preprocess_graph(original_graph) 270 entities_by_type: dict[str, list[str]] = get_entities_by_type(filtered_graph) 271 272 # Permute node labels to remove noise from the LUBM generator (e.g Professor1 always writes Paper1 etc..). 273 mapping: dict[str, str] = create_permutation(filtered_graph, entities_by_type) 274 KG: nx.MultiDiGraph = nx.relabel_nodes(filtered_graph, mapping) 275 276 path, pos, neg = sample_question(args, KG, entities_by_type) 277 start_node_pos, end_node_pos = pos 278 start_node_neg, end_node_neg = neg 279 280 def make_entry(polarity: bool, start_node: str, end_node: str) -> dict[str, Any]: 281 return { 282 'question': 'PLACEHOLDER', 283 'question_path': path, 284 'label': polarity, 285 'graph': str(graph_path), 286 'og_path': [[start_node], [end_node]], 287 'k': args.hops 288 } 289 290 json.dump(make_entry(True, start_node_pos, end_node_pos), question_file) 291 question_file.write("\n") 292 json.dump(make_entry(False, start_node_neg, end_node_neg), question_file) 293 question_file.write("\n") 294 nx.write_graphml_lxml(KG, graph_path)