errors.py (5587B)
1 import argparse 2 import datetime 3 import json 4 import pathlib 5 import sys 6 from typing import Any, Dict, Set, Optional 7 8 import datasets 9 10 11 def load_predictions(prediction_path: pathlib.Path) -> Dict[int, Dict[str, Any]]: 12 predictions = {} 13 with prediction_path.open("r") as f: 14 for line_number, line in enumerate(f): 15 obj = json.loads(line) 16 if "id" not in obj: 17 sys.exit(1) 18 predictions[obj["id"]] = obj 19 return predictions 20 21 22 def load_dataset(split_path: pathlib.Path) -> Dict[str, datasets.Dataset]: 23 splits = { 24 split: datasets.load_from_disk(str(split_path / ".." / split)) 25 for split in ["train", "validation", "test"] 26 } 27 return splits 28 29 30 def extract_actor_set(dataset: datasets.Dataset) -> Set[str]: 31 actors = set() 32 for sample in dataset: 33 if "side_a_name" in sample and sample["side_a_name"]: 34 actors.add(sample["side_a_name"]) 35 if "side_b_name" in sample and sample["side_b_name"]: 36 actors.add(sample["side_b_name"]) 37 return actors 38 39 40 def evaluate_actor_accuracy_on_unseen( 41 predictions: Dict[int, Dict[str, Any]], 42 references: datasets.Dataset, 43 train_actors: Set[str], 44 key: str, 45 ) -> float: 46 total = 0 47 correct = 0 48 for ref in references: 49 if key not in ref or ref["id"] not in predictions: 50 continue 51 gold = ref[key] 52 pred = predictions[ref["id"]].get(key) 53 if gold not in train_actors: 54 total += 1 55 if gold == pred: 56 correct += 1 57 print(f"Proportion of unseen actors in {key} predictions: {total/len(references)*100:.2f}%") 58 return correct / total if total > 0 else float("nan") 59 60 def parse_date(date_str: str) -> Optional[datetime.date]: 61 try: 62 return datetime.datetime.strptime(date_str, "%Y-%m-%d").date() 63 except (ValueError, TypeError): 64 return None 65 66 67 def evaluate_date_accuracy_on_nonzero_duration( 68 predictions: Dict[int, Dict[str, Any]], 69 references: datasets.Dataset, 70 ) -> Dict[str, float]: 71 total = 0 72 correct_start = 0 73 correct_end = 0 74 75 for ref in references: 76 if ref["id"] not in predictions: 77 continue 78 79 gold_start = ref.get("start_date", "") 80 gold_end = ref.get("end_date", "") 81 if not gold_start or not gold_end: 82 continue 83 84 if (gold_end - gold_start).days == 0: 85 continue # Skip zero-duration events 86 87 pred = predictions[ref["id"]] 88 pred_start = parse_date(pred.get("start_date", "")) 89 pred_end = parse_date(pred.get("end_date", "")) 90 if not pred_start or not pred_end: 91 continue 92 93 total += 1 94 if gold_start == pred_start: 95 correct_start += 1 96 if gold_end == pred_end: 97 correct_end += 1 98 99 return { 100 "start_date_acc_nw": round(correct_start / total, 3) if total > 0 else float("nan"), 101 "end_date_acc_nw": round(correct_end / total, 3) if total > 0 else float("nan"), 102 } 103 104 105 def evaluate_date_accuracy_on_nonzero_reporting( 106 predictions: Dict[int, Dict[str, Any]], 107 references: datasets.Dataset, 108 ) -> Dict[str, float]: 109 total = 0 110 correct_start = 0 111 correct_end = 0 112 113 for ref in references: 114 if ref["id"] not in predictions: 115 continue 116 117 gold_start = ref.get("start_date", "") 118 gold_source = ref.get("source_date", "").date() 119 if not gold_start or not gold_source: 120 continue 121 122 if (gold_source - gold_start).days == 0: 123 continue # Skip zero-reporting delay events 124 125 pred = predictions[ref["id"]] 126 pred_start = parse_date(pred.get("start_date", "")) 127 pred_end = parse_date(pred.get("end_date", "")) 128 if not pred_start or not pred_end: 129 continue 130 131 total += 1 132 if gold_start == pred_start: 133 correct_start += 1 134 if gold_source == pred_end: 135 correct_end += 1 136 137 return { 138 "start_date_acc_nr": round(correct_start / total, 3) if total > 0 else float("nan"), 139 "end_date_acc_nr": round(correct_end / total, 3) if total > 0 else float("nan"), 140 } 141 142 143 def main(split_path: pathlib.Path, prediction_path: pathlib.Path) -> None: 144 predictions = load_predictions(prediction_path) 145 splits = load_dataset(split_path) 146 147 test_split = splits[split_path.stem] 148 train_split = splits.get("train", []) 149 150 # Actor analysis 151 seen_actors = extract_actor_set(train_split) 152 153 acc_a = evaluate_actor_accuracy_on_unseen(predictions, test_split, seen_actors, "side_a_name") 154 acc_b = evaluate_actor_accuracy_on_unseen(predictions, test_split, seen_actors, "side_b_name") 155 156 print(f"Actor accuracy for side_a_name (unseen in training): {acc_a:.3f}") 157 print(f"Actor accuracy for side_b_name (unseen in training): {acc_b:.3f}") 158 159 # Date analysis 160 date_acc_nw = evaluate_date_accuracy_on_nonzero_duration(predictions, test_split) 161 print(f"Start/end date accuracies (event estimation window ≠ 0): {date_acc_nw}") 162 163 date_accs_nr = evaluate_date_accuracy_on_nonzero_reporting(predictions, test_split) 164 print(f"Start/end date accuracies (event reporting delay ≠ 0): {date_accs_nr}") 165 166 167 if __name__ == "__main__": 168 parser = argparse.ArgumentParser() 169 parser.add_argument("split", type=pathlib.Path, help="Path to the target dataset split.") 170 parser.add_argument("prediction", type=pathlib.Path, help="Path to the jsonl prediction file.") 171 args = parser.parse_args() 172 173 main(args.split, args.prediction)