Compute top-k recall

Author: Marcin Kardas

sota_extractor2/helpers/explainers.py

@@ -168,11 +168,14 @@ def _get_sota_records(self, paper):
         records.index.rename("cell_ext_id", inplace=True)
         return records
 
-    def linking_metrics(self, experiment_name="unk"):
+    def linking_metrics(self, experiment_name="unk", topk_metrics=False, filtered=True):
         paper_ids = list(self.le.proposals.keys())
 
         proposals = pd.concat(self.le.proposals.values())
-        proposals = proposals[~proposals.index.isin(self.fe.reason.index)]
+
+        # if not topk_metrics:
+        if filtered:
+            proposals = proposals[~proposals.index.isin(self.fe.reason.index)]
 
         papers = {paper_id: self.paper_collection.get_by_id(paper_id) for paper_id in paper_ids}
         missing = [paper_id for paper_id, paper in papers.items() if paper is None]
@@ -202,9 +205,10 @@ def linking_metrics(self, experiment_name="unk"):
         if "experiment_name" in df.columns:
             del df["experiment_name"]
 
-        metrics = Metrics(df, experiment_name=experiment_name)
+        metrics = Metrics(df, experiment_name=experiment_name, topk_metrics=topk_metrics)
         return metrics
 
+
     def optimize_filters(self, metrics_info):
         results = optimize_filters(self, metrics_info)
         return results

sota_extractor2/models/linking/metrics.py

@@ -14,12 +14,13 @@ class CM:
     tn: float = 0
 
 class Metrics:
-    def __init__(self, df, experiment_name="unk"):
+    def __init__(self, df, experiment_name="unk", topk_metrics=False):
         # TODO fix this, it mask the fact that our model may return more values than it should for "model
         #self.df = df[~df["model_type_gold"].str.contains('not-present') | df["model_type_pred"].str.contains('model-best')]
         self.df = df[df["model_type_gold"].str.contains('model-best') | df["model_type_pred"].str.contains('model-best')]
         self.experiment_name = experiment_name
         self.metric_type = 'best'
+        self.topk_metrics = topk_metrics
 
     def matching(self, *col_names):
         return np.all([self.df[f"{name}_pred"] == self.df[f"{name}_gold"] for name in col_names], axis=0)
@@ -42,6 +43,11 @@ def binary_confusion_matrix(self, *col_names, best_only=True):
         gold_positive = relevant_gold
         equal = self.matching(*col_names)
 
+        if self.topk_metrics:
+            equal = pd.Series(equal, index=pred_positive.index).groupby('cell_ext_id').max()
+            pred_positive = pred_positive.groupby('cell_ext_id').head(1)
+            gold_positive = gold_positive.groupby('cell_ext_id').head(1)
+
         tp = (equal & pred_positive & gold_positive).sum()
         tn = (equal & ~pred_positive & ~gold_positive).sum()
         fp = (pred_positive & (~equal | ~gold_positive)).sum()
@@ -136,4 +142,4 @@ def show(self, df):
         pd.set_option('display.max_colwidth', old_width)
 
     def show_errors(self):
-        self.show(self.errors())
+        self.show(self.errors())