|
from typing import Dict, List |
|
|
|
from lightning.pytorch.callbacks import Callback |
|
from reader.data.relik_reader_sample import RelikReaderSample |
|
|
|
from relik.reader.relik_reader_predictor import RelikReaderPredictor |
|
from relik.reader.utils.metrics import compute_metrics |
|
|
|
|
|
class StrongMatching: |
|
def __call__(self, predicted_samples: List[RelikReaderSample]) -> Dict: |
|
|
|
correct_predictions, total_predictions, total_gold = ( |
|
0, |
|
0, |
|
0, |
|
) |
|
correct_predictions_strict, total_predictions_strict = ( |
|
0, |
|
0, |
|
) |
|
correct_predictions_bound, total_predictions_bound = ( |
|
0, |
|
0, |
|
) |
|
correct_span_predictions, total_span_predictions, total_gold_spans = 0, 0, 0 |
|
|
|
|
|
for sample in predicted_samples: |
|
if sample.triplets is None: |
|
sample.triplets = [] |
|
|
|
if sample.entity_candidates: |
|
predicted_annotations_strict = set( |
|
[ |
|
( |
|
triplet["subject"]["start"], |
|
triplet["subject"]["end"], |
|
triplet["subject"]["type"], |
|
triplet["relation"]["name"], |
|
triplet["object"]["start"], |
|
triplet["object"]["end"], |
|
triplet["object"]["type"], |
|
) |
|
for triplet in sample.predicted_relations |
|
] |
|
) |
|
gold_annotations_strict = set( |
|
[ |
|
( |
|
triplet["subject"]["start"], |
|
triplet["subject"]["end"], |
|
triplet["subject"]["type"], |
|
triplet["relation"]["name"], |
|
triplet["object"]["start"], |
|
triplet["object"]["end"], |
|
triplet["object"]["type"], |
|
) |
|
for triplet in sample.triplets |
|
] |
|
) |
|
predicted_spans_strict = set(sample.predicted_entities) |
|
gold_spans_strict = set(sample.entities) |
|
|
|
correct_span_predictions += len( |
|
predicted_spans_strict.intersection(gold_spans_strict) |
|
) |
|
total_span_predictions += len(predicted_spans_strict) |
|
total_gold_spans += len(gold_spans_strict) |
|
correct_predictions_strict += len( |
|
predicted_annotations_strict.intersection(gold_annotations_strict) |
|
) |
|
total_predictions_strict += len(predicted_annotations_strict) |
|
|
|
predicted_annotations = set( |
|
[ |
|
( |
|
triplet["subject"]["start"], |
|
triplet["subject"]["end"], |
|
-1, |
|
triplet["relation"]["name"], |
|
triplet["object"]["start"], |
|
triplet["object"]["end"], |
|
-1, |
|
) |
|
for triplet in sample.predicted_relations |
|
] |
|
) |
|
gold_annotations = set( |
|
[ |
|
( |
|
triplet["subject"]["start"], |
|
triplet["subject"]["end"], |
|
-1, |
|
triplet["relation"]["name"], |
|
triplet["object"]["start"], |
|
triplet["object"]["end"], |
|
-1, |
|
) |
|
for triplet in sample.triplets |
|
] |
|
) |
|
predicted_spans = set( |
|
[(ss, se) for (ss, se, _) in sample.predicted_entities] |
|
) |
|
gold_spans = set([(ss, se) for (ss, se, _) in sample.entities]) |
|
total_gold_spans += len(gold_spans) |
|
|
|
correct_predictions_bound += len(predicted_spans.intersection(gold_spans)) |
|
total_predictions_bound += len(predicted_spans) |
|
|
|
total_predictions += len(predicted_annotations) |
|
total_gold += len(gold_annotations) |
|
|
|
correct_predictions += len( |
|
predicted_annotations.intersection(gold_annotations) |
|
) |
|
|
|
span_precision, span_recall, span_f1 = compute_metrics( |
|
correct_span_predictions, total_span_predictions, total_gold_spans |
|
) |
|
bound_precision, bound_recall, bound_f1 = compute_metrics( |
|
correct_predictions_bound, total_predictions_bound, total_gold_spans |
|
) |
|
|
|
precision, recall, f1 = compute_metrics( |
|
correct_predictions, total_predictions, total_gold |
|
) |
|
|
|
if sample.entity_candidates: |
|
precision_strict, recall_strict, f1_strict = compute_metrics( |
|
correct_predictions_strict, total_predictions_strict, total_gold |
|
) |
|
return { |
|
"span-precision": span_precision, |
|
"span-recall": span_recall, |
|
"span-f1": span_f1, |
|
"precision": precision, |
|
"recall": recall, |
|
"f1": f1, |
|
"precision-strict": precision_strict, |
|
"recall-strict": recall_strict, |
|
"f1-strict": f1_strict, |
|
} |
|
else: |
|
return { |
|
"span-precision": bound_precision, |
|
"span-recall": bound_recall, |
|
"span-f1": bound_f1, |
|
"precision": precision, |
|
"recall": recall, |
|
"f1": f1, |
|
} |
|
|
|
|
|
class REStrongMatchingCallback(Callback): |
|
def __init__(self, dataset_path: str, dataset_conf) -> None: |
|
super().__init__() |
|
self.dataset_path = dataset_path |
|
self.dataset_conf = dataset_conf |
|
self.strong_matching_metric = StrongMatching() |
|
|
|
def on_validation_epoch_start(self, trainer, pl_module) -> None: |
|
relik_reader_predictor = RelikReaderPredictor(pl_module.relik_reader_re_model) |
|
predicted_samples = relik_reader_predictor._predict( |
|
self.dataset_path, |
|
None, |
|
self.dataset_conf, |
|
) |
|
predicted_samples = list(predicted_samples) |
|
for k, v in self.strong_matching_metric(predicted_samples).items(): |
|
pl_module.log(f"val_{k}", v) |
|
|
