dataset_maker.py (10917B)
1 from pathlib import Path 2 from tqdm import tqdm 3 from typing import Dict, List, Any 4 import argparse 5 import copy 6 import json 7 import jsonlines 8 import logging 9 import math 10 import networkx as nx 11 import random 12 13 14 POSSIBLE_TASKS = ["sub", "pro", "sys"] 15 SUB_TRAINING_SIZE = 10000 16 SUB_TEST_SIZE = 2000 17 ANNOTAITON_MAP_PATH = Path("./static/annotation_map.json") 18 HOP_2_ANNOTATIONS = Path("./static/2_hop_annotations.json") 19 HOP_3_ANNOTATIONS = Path("./static/3_hop_annotations.json") 20 HOP_4_ANNOTATIONS = Path("./static/4_hop_annotations.json") 21 HOP_2_PATH = Path("./data/2_hop_questions/2_hop_qs_25000.jsonl") 22 HOP_3_PATH = Path("./data/3_hop_questions/3_hop_qs_25000.jsonl") 23 HOP_4_PATH = Path("./data/4_hop_questions/4_hop_qs_25000.jsonl") 24 25 26 def chunking(lst, n): 27 for i in range(0, len(lst), n): 28 yield lst[i:i + n] 29 30 31 def parse_args() -> argparse.Namespace: 32 parser = argparse.ArgumentParser() 33 parser.add_argument("--task", type=str, default="sub", help=f"Which task to produce dataset splits for (one of: {' '.join(POSSIBLE_TASKS)}).") 34 parser.add_argument("--hop_2_path", type=Path, default=HOP_2_PATH) 35 parser.add_argument("--hop_3_path", type=Path, default=HOP_3_PATH) 36 parser.add_argument("--hop_4_path", type=Path, default=HOP_4_PATH) 37 parser.add_argument("--shuffle", action="store_true") 38 parser.add_argument("--round", type=int, help="what round to execute for the sys task") 39 return parser.parse_args() 40 41 42 def get_json(path: Path) -> Dict: 43 with open(path, 'r') as f: 44 amap = json.load(f) 45 return amap 46 47 48 def get_jsonl(path: Path) -> List: 49 with jsonlines.open(path, 'r') as f: 50 lines = [line for line in f] 51 return lines 52 53 54 def write_jsonl(path: Path, data: List) -> None: 55 path.parent.mkdir(exist_ok=True, parents=True) 56 with jsonlines.open(path, 'w') as f: 57 f.write_all(data) 58 logging.info(f"Wrote {path} to disk...") 59 60 61 def get_annotation_map(path: Path): 62 index_2_relation = get_json(path) 63 relations_2_id = {} 64 for relation, value in index_2_relation.items(): 65 relations_2_id[tuple(value['type'])] = relation 66 return index_2_relation, relations_2_id 67 68 69 def create_individual_hop_files(hop_path: Path, index_2_relation: Dict, relations_2_id: Dict): 70 raw_qs = get_jsonl(hop_path) 71 questions = [] 72 errors = 0 73 lacking_annotation = 0 74 unseen_relation_types = set() 75 for question in raw_qs: 76 if question['og_path'][0][0] is None: 77 errors += 1 78 continue 79 idx = relations_2_id.get(tuple(question['question_path']), -1) 80 if idx == -1: 81 unseen_relation_types.add(tuple(question['question_path'])) 82 lacking_annotation += 1 83 continue 84 annotation = index_2_relation[idx]['mapping'] 85 q = annotation.replace('X', question['og_path'][0][0]) 86 q = q.replace('Y', question['og_path'][-1][-1]) 87 question['question'] = q 88 questions.append(question) 89 return questions, errors, lacking_annotation, unseen_relation_types 90 91 92 def assert_primitives(train_relations: List, test_relations: List) -> bool: 93 train_primitives = set() 94 test_primitives = set() 95 for r in train_relations: 96 for t in r: 97 train_primitives.add(t) 98 for r in test_relations: 99 for t in r: 100 test_primitives.add(t) 101 good_2_go = True 102 for p in test_primitives: 103 if p not in train_primitives: 104 good_2_go = False 105 return good_2_go 106 107 108 def load_graph_from_path(path: str) -> nx.Graph: 109 G = nx.read_graphml(f'{path}') 110 return G 111 112 113 def get_spelled_out_path(KG, list_of_vertices, only_relations): 114 """ Transform a list of vertices into a list of (head, relation, tail) or a list relations. """ 115 path: list[str | list[str]] = [] 116 for i in range(len(list_of_vertices)-1): 117 head: str = list_of_vertices[i] 118 tail: str = list_of_vertices[i+1] 119 kg_edge: dict[str, Any] = KG.get_edge_data(head, tail) 120 predicate: str 121 if isinstance(KG, nx.MultiDiGraph): 122 predicate = list(kg_edge.keys())[0] # no parallel edges going in the same direction. 123 else: # nx.DiGraph 124 predicate = kg_edge["type"] 125 path.append(predicate if only_relations else [head, predicate, tail]) 126 return path 127 128 129 def prune_data(data_points: List[List], test_relations: List, hop: int) -> List: 130 pruned_samples = [] 131 for point in tqdm(data_points): 132 graph_path = point["graph"] 133 KG = load_graph_from_path(graph_path) 134 head_p = point["og_path"][0][0] 135 tail_p = point["og_path"][-1][-1] 136 paths = [n for n in nx.all_simple_paths(KG, head_p, tail_p, hop+1) if len(n) == hop+1] 137 to_break = False 138 for path in paths: 139 path_key = tuple(get_spelled_out_path(KG, path, True)) 140 if path_key in test_relations: 141 to_break = True 142 continue 143 if not to_break: 144 pruned_samples.append(point) 145 return pruned_samples 146 147 148 if __name__ == "__main__": 149 logging.basicConfig(level=logging.INFO) 150 logging.info("Dataset creator") 151 args = parse_args() 152 assert args.task in POSSIBLE_TASKS, exit(f"Invalid task argument provided, use one of: {' '.join(POSSIBLE_TASKS)}") 153 index_2_relation, relations_2_id = get_annotation_map(ANNOTAITON_MAP_PATH) 154 155 hop_2_questions, hop_2_errors, hop_2_lacking_annotation, hop_2_unseen_relation_types = create_individual_hop_files(args.hop_2_path, index_2_relation, relations_2_id) 156 logging.info(f"Parsed 2 hops, total question: {len(hop_2_questions)}, errors: {hop_2_errors}, missing annotation: {hop_2_lacking_annotation}") 157 158 hop_3_questions, hop_3_errors, hop_3_lacking_annotation, hop_3_unseen_relation_types = create_individual_hop_files(args.hop_3_path, index_2_relation, relations_2_id) 159 logging.info(f"Parsed 3 hops, total question: {len(hop_3_questions)}, errors: {hop_3_errors}, missing annotation: {hop_3_lacking_annotation}") 160 161 hop_4_questions, hop_4_errors, hop_4_lacking_annotation, hop_4_unseen_relation_types = create_individual_hop_files(args.hop_4_path, index_2_relation, relations_2_id) 162 logging.info(f"Parsed 4 hops, total question: {len(hop_4_questions)}, errors: {hop_4_errors}, missing annotation: {hop_4_lacking_annotation}") 163 164 # Creating dataset splits 165 if args.task == "sub": 166 save_path: Path = Path("./data/substitutivity") 167 for hop, qs in zip(["2", "3", "4"], [hop_2_questions, hop_3_questions, hop_4_questions]): 168 if args.shuffle: 169 random.shuffle(qs) 170 train_data = qs[:SUB_TRAINING_SIZE] 171 test_data = qs[SUB_TRAINING_SIZE:SUB_TRAINING_SIZE+SUB_TEST_SIZE] 172 train_path = save_path / f"{hop}_hop_train.jsonl" 173 test_path = save_path / f"{hop}_hop_test.jsonl" 174 write_jsonl(train_path, train_data) 175 write_jsonl(test_path, test_data) 176 elif args.task == "pro": 177 save_path: Path = Path("./data/productivity") 178 if args.shuffle: 179 random.shuffle(hop_2_questions) 180 random.shuffle(hop_3_questions) 181 random.shuffle(hop_4_questions) 182 train_2 = hop_2_questions[:5000] 183 train_3 = hop_3_questions[:5000] 184 train_4 = hop_4_questions[:5000] 185 test_2 = hop_2_questions[-2000:] 186 test_3 = hop_3_questions[-2000:] 187 test_4 = hop_4_questions[-2000:] 188 train_2_and_3 = train_2 + train_3 189 train_2_and_4 = train_2 + train_4 190 train_3_and_4 = train_3 + train_4 191 train_2_and_3_path = save_path / "2and3_train.jsonl" 192 train_2_and_4_path = save_path / "2and4_train.jsonl" 193 train_3_and_4_path = save_path / "3and4_train.jsonl" 194 test_2_path = save_path / "2_test.jsonl" 195 test_3_path = save_path / "3_test.jsonl" 196 test_4_path = save_path / "4_test.jsonl" 197 write_jsonl(train_2_and_3_path, train_2_and_3) 198 write_jsonl(train_2_and_4_path, train_2_and_4) 199 write_jsonl(train_3_and_4_path, train_3_and_4) 200 write_jsonl(test_2_path, test_2) 201 write_jsonl(test_3_path, test_3) 202 write_jsonl(test_4_path, test_4) 203 elif args.task == "sys": 204 print(f"Creating data for round {i_round}") 205 _, hop2_rel_to_id = get_annotation_map(HOP_2_ANNOTATIONS) 206 _, hop3_rel_to_id = get_annotation_map(HOP_3_ANNOTATIONS) 207 _, hop4_rel_to_id = get_annotation_map(HOP_4_ANNOTATIONS) 208 save_path: Path = Path("./data/systematicity") 209 hop_to_data = { 210 2: hop_2_questions, 211 3: hop_3_questions, 212 4: hop_4_questions, 213 } 214 n_chunks = 5 215 hop_to_percentage = { 216 2: 0.3, 217 3: 0.2, 218 4: 0.1 219 } 220 for hop, hop_map in zip([2, 3, 4], [hop2_rel_to_id, hop3_rel_to_id, hop4_rel_to_id]): 221 has_valid_setup = False 222 tries = 10 223 while not has_valid_setup: 224 if tries == 0: 225 exit(f"Failed to create split for {hop}-hop setting") 226 possible_relations = list(hop_map.keys()) 227 sample_size = math.ceil(hop_to_percentage[hop] * len(possible_relations)) 228 test_relations = random.sample(possible_relations, sample_size) 229 train_relations = [r for r in possible_relations if r not in test_relations] 230 has_valid_setup = assert_primitives(train_relations, test_relations) 231 tries -= 1 232 233 test_data = [] 234 tmp_train_data = [] 235 rouge_points = 0 236 for data_point in hop_to_data[hop]: 237 relation = tuple(data_point["question_path"]) 238 if relation in train_relations: 239 tmp_train_data.append(data_point) 240 elif relation in test_relations: 241 test_data.append(data_point) 242 else: 243 rouge_points += 1 244 245 logging.info(f"Found rouge data points count: {rouge_points}") 246 logging.info(f"Total train size: {len(tmp_train_data)}") 247 logging.info(f"Total test size: {len(test_data)}") 248 pruned_train_data = prune_data(tmp_train_data, test_relations, hop) 249 logging.info(f"Pruned {1 - (len(pruned_train_data) / len(tmp_train_data))}% for {hop}-hop") 250 chunks = list(chunking(pruned_train_data, len(pruned_train_data) // n_chunks)) 251 logging.info(f"N-chunks: {len(chunks)}") 252 for i in range(0, len(chunks)-1): 253 td = copy.deepcopy(chunks) 254 popped = td.pop(i) 255 td = [j for sub in td for j in sub] 256 output_path = save_path / f"round_{i_round}_{hop}hop_fold{i}_train.jsonl" 257 write_jsonl(output_path, td) 258 test_output_path = save_path / f"round_{i_round}_{hop}_test.jsonl" 259 write_jsonl(test_output_path, test_data)