add mention detection metrics reporting for train/test scripts

Author: Aethor

tibert/run_test.py

@@ -15,7 +15,7 @@
     BertForCoreferenceResolution,
     CamembertForCoreferenceResolution,
 )
-from tibert.score import score_coref_predictions
+from tibert.score import score_coref_predictions, score_mention_detection
 from tibert.predict import predict_coref
 from tibert.utils import split_coreference_document_tokens
 
@@ -82,7 +82,7 @@ def main(
     all_annotated_docs = []
     for document in tqdm(test_dataset.documents):
         doc_dataset = CoreferenceDataset(
-            split_coreference_document_tokens(document, 512),
+            [document],
             tokenizer,
             max_span_size,
         )
@@ -110,6 +110,17 @@ def main(
             annotated_doc = CoreferenceDocument.concatenated(annotated_docs)
         all_annotated_docs.append(annotated_doc)
 
+    mention_pre, mention_rec, mention_f1 = score_mention_detection(
+        all_annotated_docs, test_dataset.documents
+    )
+    for metric_key, score in [
+        ("precision", mention_pre),
+        ("recall", mention_rec),
+        ("f1", mention_f1),
+    ]:
+        print(f"mention.{metric_key}={score}")
+        _run.log_scalar(f"mention.{metric_key}", score)
+
     scores = score_coref_predictions(all_annotated_docs, test_dataset.documents)
     for key, score_dict in scores.items():
         for metric_key, score in score_dict.items():

tibert/run_train.py

@@ -14,6 +14,7 @@
     load_train_checkpoint,
     predict_coref,
     score_coref_predictions,
+    score_mention_detection,
 )
 from tibert.bertcoref import CoreferenceDataset, load_democrat_dataset
 
@@ -145,6 +146,17 @@ def main(
     )
     assert isinstance(annotated_docs, list)
 
+    mention_pre, mention_rec, mention_f1 = score_mention_detection(
+        annotated_docs, test_dataset.documents
+    )
+    for metric_key, score in [
+        ("precision", mention_pre),
+        ("recall", mention_rec),
+        ("f1", mention_f1),
+    ]:
+        print(f"mention.{metric_key}={score}")
+        _run.log_scalar(f"mention.{metric_key}", score)
+
     metrics = score_coref_predictions(annotated_docs, test_dataset.documents)
     print(metrics)
     for key, score_dict in metrics.items():