evaluate.py (11273B)
1 from __future__ import annotations 2 3 from typing import Any, ClassVar, cast 4 import argparse 5 import datetime 6 import itertools 7 import json 8 import logging 9 import pathlib 10 import sys 11 12 import datasets 13 import tqdm 14 import torch 15 import transformers 16 17 18 logger: logging.Logger = logging.getLogger(__name__) 19 20 21 FIELDS: set[str] = { 22 "side_a_name", "side_b_name", # Actors 23 "start_date", "end_date", # Time 24 "location_root_name", "location_adm1_name", "location_adm2_name", "location_where_name", # Location 25 "deaths_side_a", "deaths_side_b", "deaths_civilian", "deaths_unknown", "deaths_low", "deaths_high", # Deaths 26 } 27 28 29 SEMANTIC_MODEL: str = "google-bert/bert-base-uncased" # "sentence-transformers/all-mpnet-base-v2" 30 31 32 class Metrics: 33 ACCURACY_FIELDS: set[str] = FIELDS 34 35 RMSE_FIELDS: set[str] = { 36 "start_date", "end_date", 37 "deaths_side_a", "deaths_side_b", "deaths_civilian", "deaths_unknown", "deaths_low", "deaths_high", 38 } 39 40 def __init__(self: Metrics, dataset: datasets.Dataset | None) -> None: 41 self.sample_count: int = 0 42 self.correct: dict[str, int] = { field: 0 for field in self.ACCURACY_FIELDS } 43 self.square_error: dict[str, int] = { field: 0 for field in self.RMSE_FIELDS } 44 self.semantic_evaluation: bool = dataset is not None 45 if self.semantic_evaluation: 46 self.cosine: dict[str, float] = { field: 0 for field in self.ACCURACY_FIELDS } 47 self.reciprocal_rank: dict[str, float] = { field: 0 for field in self.ACCURACY_FIELDS } 48 self.precision1: dict[str, int] = { field: 0 for field in self.ACCURACY_FIELDS } 49 self._semantic_init(dataset) 50 51 def _semantic_init(self: Metrics, dataset: datasets.Dataset) -> None: 52 self._semantic_tokenizer = transformers.AutoTokenizer.from_pretrained(SEMANTIC_MODEL) 53 self._semantic_model = transformers.AutoModel.from_pretrained(SEMANTIC_MODEL) 54 self._candidates: dict[str, list[str]] = { field: list(map(str, set(dataset[field]))) for field in self.ACCURACY_FIELDS } 55 self._embedding_cache: dict[str, torch.Tensor] = {} 56 all_values: set[str] = set(itertools.chain(*self._candidates.values())) 57 self._embedding_cache = {key: self._semantic_embedding(key) for key in tqdm.tqdm(all_values, desc="Embedding gold candidates")} 58 self._candidates_embeddings: dict[str, torch.Tensor] = {field: torch.stack([self._embedding_cache[key] for key in candidate_values]) for field, candidate_values in self._candidates.items()} 59 60 def _semantic_embedding(self: Metrics, value: str) -> torch.Tensor: 61 if value in self._embedding_cache: 62 return self._embedding_cache[value] 63 with torch.no_grad(): 64 input_tensors: dict[str, torch.Tensor] = self._semantic_tokenizer(value, return_tensors="pt") 65 cls_embedding: torch.Tensor = self._semantic_model(**input_tensors).last_hidden_state[0,0] 66 return cls_embedding 67 68 def _update_semantic_evaluation(self: Metrics, field: str, gold: str, prediction: str | None) -> None: 69 if prediction is None: 70 self.cosine[field] += -1 71 return # MRR is 0 72 73 gold_index: int = self._candidates[field].index(gold) 74 prediction_embedding: torch.Tensor = self._semantic_embedding(prediction) 75 similarities: torch.Tensor = torch.cosine_similarity(self._candidates_embeddings[field], prediction_embedding) 76 self.cosine[field] += cast(float, similarities[gold_index].item()) 77 ranks: list[int] = cast(list[int], torch.argsort(similarities, descending=True).tolist()) 78 gold_rank: int = ranks.index(gold_index) 79 self.reciprocal_rank[field] += 1/(1+gold_rank) 80 self.precision1[field] += (gold_rank == 0) 81 82 def update(self: Metrics, gold: dict[str, Any], prediction: dict[str, Any]) -> None: 83 self.sample_count += 1 84 85 missing_fields: set[str] = FIELDS - prediction.keys() 86 if missing_fields: 87 logger.warning(f"Missing predicted field{'s' if len(missing_fields)>1 else ''} for sample id={gold['id']}: {', '.join(missing_fields)}") 88 89 gold = {key: value.strip() if isinstance(value, str) else value for key, value in gold.items()} 90 prediction = {key: value.strip() if isinstance(value, str) else value for key, value in prediction.items()} 91 92 # The actors are permutation invariant, if inverting them would get a better accuracy 93 if gold["side_a_name"] != prediction.get("side_a_name") and gold["side_b_name"] != prediction.get("side_b_name") \ 94 and (gold["side_a_name"] == prediction.get("side_b_name") or gold["side_b_name"] == prediction.get("side_a_name")): 95 logger.debug(f"Switching sides for sample id={gold['id']}") 96 # invert the name of the actors 97 prediction["side_a_name"], prediction["side_b_name"] = prediction.get("side_b_name"), prediction.get("side_a_name") 98 # invert the number of deaths on each side 99 prediction["deaths_side_a"], prediction["deaths_side_b"] = prediction.get("deaths_side_b"), prediction.get("deaths_side_a") 100 101 for field in self.ACCURACY_FIELDS: 102 self.correct[field] += (str(prediction.get(field, "")) == str(gold[field])) 103 if self.semantic_evaluation: 104 self._update_semantic_evaluation(field, str(gold[field]), str(prediction[field]) if field in prediction else None) 105 106 for field in self.RMSE_FIELDS: 107 # The field might have been set to None by the actor inversion logic 108 prediction_value = prediction.get(field) 109 if field.startswith("deaths_"): 110 if prediction_value is None: 111 prediction_value = 0 112 difference = (gold[field] - prediction_value)**2 113 elif field.endswith("_date"): 114 if isinstance(prediction_value, str): 115 try: 116 prediction_value = datetime.date.fromisoformat(prediction_value) 117 except ValueError: 118 logger.info(f"Couldn't parse {field} as date for sample id={gold['id']}") 119 prediction_value = None 120 if prediction_value is None: 121 prediction_value = gold["source_date"].date() 122 difference = (gold[field] - prediction_value).days**2 123 self.square_error[field] += difference 124 125 def report(self: Metrics) -> None: 126 measures: list[str] = ["accuracy"] 127 atoms: dict[str, float] = {} 128 for field in self.ACCURACY_FIELDS: 129 atoms[f"{field} accuracy"] = 100 * self.correct[field] / self.sample_count 130 for field in self.RMSE_FIELDS: 131 atoms[f"{field} RMSE"] = (self.square_error[field] / self.sample_count)**0.5 132 if self.semantic_evaluation: 133 measures.extend(["semantic MRR", "semantic p@1", "semantic cos"]) 134 for field in self.ACCURACY_FIELDS: 135 atoms[f"{field} semantic MRR"] = self.reciprocal_rank[field] / self.sample_count 136 atoms[f"{field} semantic p@1"] = self.precision1[field] / self.sample_count 137 atoms[f"{field} semantic cos"] = self.cosine[field] / self.sample_count 138 139 aggregates: dict[str, float] = {} 140 for measure in measures: 141 aggregates[f"actor {measure}"] = (atoms[f"side_a_name {measure}"] + atoms[f"side_b_name {measure}"]) / 2 142 aggregates[f"date {measure}"] = (atoms[f"start_date {measure}"] + atoms[f"end_date {measure}"]) / 2 143 aggregates[f"location {measure}"] = (atoms[f"location_root_name {measure}"] + atoms[f"location_adm1_name {measure}"] + atoms[f"location_adm2_name {measure}"] + atoms[f"location_where_name {measure}"]) / 4 144 aggregates[f"deaths {measure}"] = (atoms[f"deaths_side_a {measure}"] + atoms[f"deaths_side_b {measure}"] + atoms[f"deaths_civilian {measure}"] + atoms[f"deaths_unknown {measure}"] + atoms[f"deaths_low {measure}"] + atoms[f"deaths_high {measure}"]) / 6 145 aggregates[f"aggregate {measure}"] = (aggregates[f"actor {measure}"] + aggregates[f"date {measure}"] + aggregates[f"location {measure}"] + aggregates[f"deaths {measure}"]) / 4 146 aggregates["date RMSE"] = (atoms["start_date RMSE"] + atoms["end_date RMSE"]) / 2 147 aggregates["deaths RMSE"] = (atoms["deaths_side_a RMSE"] + atoms["deaths_side_b RMSE"] + atoms["deaths_civilian RMSE"] + atoms["deaths_unknown RMSE"] + atoms["deaths_low RMSE"] + atoms["deaths_high RMSE"]) / 6 148 149 print(json.dumps(atoms | aggregates)) 150 for result in [atoms, aggregates]: 151 print(38*"=") 152 for field, value in result.items(): 153 if field.endswith("accuracy"): 154 print(f"{field:<32} | {value:7.1f}") 155 else: 156 print(f"{field:<32} | {value:7.4f}") 157 158 159 def parse_command_argument() -> argparse.Namespace: 160 """ Parse command line arguments. """ 161 parser = argparse.ArgumentParser(description="Update a huggingface dataset using a minhash input and output.") 162 parser.add_argument("split", type=pathlib.Path, help="Path to the target dataset split.") 163 parser.add_argument("prediction", type=pathlib.Path, help="Path to the jsonl prediction file.") 164 parser.add_argument("-d", "--debug", action="store_const", dest="loglevel", const=logging.DEBUG, default=logging.WARNING, help="Print lots of debugging statements.") 165 parser.add_argument("-s", "--no-semantic", action="store_true", help="Do not run semantic evaluation.") 166 parser.add_argument("-v", "--verbose", action="store_const", dest="loglevel", const=logging.INFO, help="Be verbose.") 167 args: argparse.Namespace = parser.parse_args() 168 return args 169 170 171 def main(split_path: pathlib.Path, prediction_path: pathlib.Path, semantic: bool) -> None: 172 predictions: dict[int, dict[str, Any]] = {} 173 with prediction_path.open("r") as prediction_file: 174 for line_number, prediction_line in enumerate(prediction_file): 175 prediction: dict[str, Any] = json.loads(prediction_line) 176 if "id" not in prediction: 177 logger.fatal(f"Prediction file contains an object without id line {line_number+1}.") 178 sys.exit(1) 179 predictions[prediction["id"]] = prediction 180 181 dataset: datasets.Dataset 182 if semantic: 183 splits: dict[str, datasets.Dataset] = {split: datasets.load_from_disk(str(split_path / ".." / split)) 184 for split in ["train", "validation", "test"]} 185 dataset = splits[split_path.stem] 186 metrics = Metrics(datasets.concatenate_datasets(splits.values())) 187 else: 188 dataset = datasets.load_from_disk(str(split_path)) 189 metrics = Metrics(None) 190 for sample in tqdm.tqdm(dataset, desc="Scoring samples"): 191 if sample["id"] not in predictions: 192 logger.warning(f"Missing prediction for sample id={sample['id']}") 193 metrics.update(sample, predictions.get(sample["id"], {})) 194 metrics.report() 195 196 197 if __name__ == "__main__": 198 args: argparse.Namespace = parse_command_argument() 199 logging.basicConfig(level=args.loglevel) 200 main(split_path=args.split, prediction_path=args.prediction, semantic=not args.no_semantic)