Introduce explainers

* implement Table deepcopy so we can store Tables in various stages of pipeline
* add explainers for:
  + table type
  + filtering
* add SessionRecorder to speed up testing
* modify ULMFiT interpretation classes to show signed words contributions

Author: mkardas

sota_extractor2/data/paper_collection.py

@@ -28,6 +28,12 @@ def __init__(self, paper_id, text, tables, annotations):
         else:
             self.gold_tags = ''
 
+    def table_by_name(self, name):
+        for table in self.tables:
+            if table.name == name:
+                return table
+        return None
+
 
 # todo: make sure multithreading/processing won't cause collisions
 def random_id():

sota_extractor2/data/table.py

@@ -6,6 +6,8 @@
 from dataclasses import dataclass, field
 from typing import List
 from ..helpers.jupyter import display_table
+from copy import deepcopy
+
 
 @dataclass
 class Cell:
@@ -68,15 +70,18 @@ def read_str_csv(filename):
     return df
 
 
+class CellDataFrame(pd.DataFrame):
+    """We subclass pandas DataFrame in order to make deepcopy recursively copy cells"""
+    def __deepcopy__(self, memodict={}):
+        return CellDataFrame(self.applymap(lambda cell: deepcopy(cell, memodict)))
 
 
 class Table:
     def __init__(self, name, df, layout, caption=None, figure_id=None, annotations=None, migrate=False, old_name=None, guessed_tags=None):
         self.name = name
-        self.df = df
         self.caption = caption
         self.figure_id = figure_id
-        self.df = df.applymap(str2cell)
+        self.df = CellDataFrame(df.applymap(str2cell))
 
         if migrate:
             self.old_name = old_name

sota_extractor2/helpers/explainers.py

@@ -0,0 +1,43 @@
+from ..models.structure import TableType
+from ..loggers import StructurePredictionEvaluator, LinkerEvaluator, FilteringEvaluator
+import pandas as pd
+
+
+class TableTypeExplainer:
+    def __init__(self, paper, table, table_type, probs):
+        self.paper = paper
+        self.table = table
+        self.table_type = table_type
+        self.probs = pd.DataFrame(probs, columns=["type", "probability"])
+
+    def __str__(self):
+        return f"Table {self.table.name} was labelled as {self.table_type}."
+
+    def display(self):
+        print(self)
+        self.probs.display()
+
+
+class Explainer:
+    def __init__(self, pipeline_logger, paper_collection):
+        self.spe = StructurePredictionEvaluator(pipeline_logger, paper_collection)
+        self.le = LinkerEvaluator(pipeline_logger, paper_collection)
+        self.fe = FilteringEvaluator(pipeline_logger)
+
+    def explain(self, paper, cell_ext_id):
+        paper_id, table_name, rc = cell_ext_id.split('/')
+        if paper.paper_id != paper_id:
+            return "No such cell"
+
+        row, col = [int(x) for x in rc.split('.')]
+
+        table_type, probs = self.spe.get_table_type_predictions(paper_id, table_name)
+
+        if table_type == TableType.IRRELEVANT:
+            return TableTypeExplainer(paper, paper.table_by_name(table_name), table_type, probs)
+
+        reason = self.fe.reason.get(cell_ext_id)
+        if reason is None:
+            pass
+        else:
+            return reason

sota_extractor2/helpers/interpret.py

@@ -0,0 +1,59 @@
+from fastai.text.interpret import TextClassificationInterpretation as AbsTextClassificationInterpretation, _eval_dropouts
+from fastai.basic_data import DatasetType
+import torch
+
+
+__all__ = ["TextClassificationInterpretation", "TextMultiClassificationInterpretation"]
+
+
+class TextClassificationInterpretation(AbsTextClassificationInterpretation):
+    @classmethod
+    def from_learner(cls, learner):
+        empty_preds = torch.Tensor([[1]])
+        return cls(learner, empty_preds, None, None)
+
+    def intrinsic_attention(self, text:str, class_id:int=None):
+        """Calculate the intrinsic attention of the input w.r.t to an output `class_id`, or the classification given by the model if `None`.
+        Similar as in base class, but does not apply abs() before summing gradients.
+        """
+        self.model.train()
+        _eval_dropouts(self.model)
+        self.model.zero_grad()
+        self.model.reset()
+        ids = self.data.one_item(text)[0]
+        emb = self.model[0].module.encoder(ids).detach().requires_grad_(True)
+        lstm_output = self.model[0].module(emb, from_embeddings=True)
+        self.model.eval()
+        cl = self.model[1](lstm_output + (torch.zeros_like(ids).byte(),))[0].softmax(dim=-1)
+        if class_id is None: class_id = cl.argmax()
+        cl[0][class_id].backward()
+        # attn = emb.grad.squeeze().abs().sum(dim=-1)
+        # attn /= attn.max()
+        attn = emb.grad.squeeze().sum(dim=-1)
+        attn = attn / attn.abs().max() * 0.5 + 0.5
+        tokens = self.data.single_ds.reconstruct(ids[0])
+        return tokens, attn
+
+
+class TextMultiClassificationInterpretation(TextClassificationInterpretation):
+    def intrinsic_attention(self, text:str, class_id:int=None):
+        """Calculate the intrinsic attention of the input w.r.t to an output `class_id`, or the classification given by the model if `None`.
+        Similar as in base class, but uses sigmoid instead of softmax and does not apply abs() before summing gradients.
+        """
+        self.model.train()
+        _eval_dropouts(self.model)
+        self.model.zero_grad()
+        self.model.reset()
+        ids = self.data.one_item(text)[0]
+        emb = self.model[0].module.encoder(ids).detach().requires_grad_(True)
+        lstm_output = self.model[0].module(emb, from_embeddings=True)
+        self.model.eval()
+        cl = self.model[1](lstm_output + (torch.zeros_like(ids).byte(),))[0].sigmoid()
+        if class_id is None: class_id = cl.argmax()
+        cl[0][class_id].backward()
+        # attn = emb.grad.squeeze().abs().sum(dim=-1)
+        # attn /= attn.max()
+        attn = emb.grad.squeeze().sum(dim=-1)
+        attn = attn / attn.abs().max() * 0.5 + 0.5
+        tokens = self.data.single_ds.reconstruct(ids[0])
+        return tokens, attn

sota_extractor2/loggers.py

@@ -2,6 +2,8 @@
 import pandas as pd
 from .models.structure.experiment import Experiment, label_map, Labels
 from .models.structure.type_predictor import TableType
+from copy import deepcopy
+import pickle
 
 
 class BaseLogger:
@@ -21,30 +23,85 @@ def __call__(self, step, **kwargs):
         print(f"[STEP] {step}: {kwargs}", file=self.file)
 
 
+class SessionRecorder:
+    def __init__(self, pipeline_logger):
+        self.pipeline_logger = pipeline_logger
+        self.session = []
+        self._recording = False
+
+    def __call__(self, step, **kwargs):
+        self.session.append((step, deepcopy(kwargs)))
+
+    def reset(self):
+        self.session = []
+
+    def record(self):
+        if not self._recording:
+            self.pipeline_logger.register(".*", self)
+            self._recording = True
+
+    def stop(self):
+        if self._recording:
+            self.pipeline_logger.unregister(".*", self)
+            self._recording = False
+
+    def replay(self):
+        self.stop()
+        for step, kwargs in self.session:
+            self.pipeline_logger(step, **kwargs)
+
+    def save_session(self, path):
+        with open(path, "wb") as f:
+            pickle.dump(self.session, f)
+
+    def load_session(self, path):
+        with open(path, "rb") as f:
+            self.session = pickle.load(f)
+
+
 class StructurePredictionEvaluator:
     def __init__(self, pipeline_logger, pc):
-        pipeline_logger.register("structure_prediction::tables_labelled", self.on_tables_labelled)
+        pipeline_logger.register("structure_prediction::evidences_split", self.on_evidences_split)
+        pipeline_logger.register("structure_prediction::tables_labeled", self.on_tables_labeled)
         pipeline_logger.register("type_prediction::predicted", self.on_type_predicted)
+        pipeline_logger.register("type_prediction::multiclass_predicted", self.on_type_multiclass_predicted)
         self.pc = pc
         self.results = {}
         self.type_predictions = {}
+        self.type_multiclass_predictions = {}
+        self.evidences = pd.DataFrame()
+
+    def on_type_multiclass_predicted(self, step, paper, tables, threshold, predictions):
+        for table, prediction in zip(tables, predictions):
+            self.type_multiclass_predictions[paper.paper_id, table.name] = {
+                TableType.SOTA: prediction[0],
+                TableType.ABLATION: prediction[1],
+                TableType.IRRELEVANT: threshold
+            }
 
     def on_type_predicted(self, step, paper, tables, predictions):
-        self.type_predictions[paper.paper_id] = predictions
+        for table, prediction in zip(tables, predictions):
+            self.type_predictions[paper.paper_id, table.name] = prediction
+
+    def on_evidences_split(self, step, evidences, evidences_num):
+        self.evidences = pd.concat([self.evidences, evidences])
 
-    def on_tables_labelled(self, step, paper, tables):
+    def on_tables_labeled(self, step, paper, labeled_tables):
         golds = [p for p in self.pc if p.text.title == paper.text.title]
         paper_id = paper.paper_id
         type_results = []
         cells_results = []
+        labeled_tables = {table.name: table for table in labeled_tables}
         if len(golds) == 1:
             gold = golds[0]
-            for gold_table, table, table_type in zip(gold.tables, paper.tables, self.type_predictions.get(paper.paper_id, [])):
+            for gold_table, table, in zip(gold.tables, paper.tables):
+                table_type = self.type_predictions[paper.paper_id, table.name]
                 is_important = table_type == TableType.SOTA or table_type == TableType.ABLATION
                 gold_is_important = "sota" in gold_table.gold_tags or "ablation" in gold_table.gold_tags
                 type_results.append({"predicted": is_important, "gold": gold_is_important, "name": table.name})
                 if not is_important:
                     continue
+                table = labeled_tables[table.name]
                 rows, cols = table.df.shape
                 for r in range(rows):
                     for c in range(cols):
@@ -76,6 +133,14 @@ def metrics(self, paper_id):
         e._set_results(paper_id, self.map_tags(results['cells'].predicted), self.map_tags(results['cells'].gold))
         e.show_results(paper_id, normalize=True)
 
+    def get_table_type_predictions(self, paper_id, table_name):
+        prediction = self.type_predictions.get((paper_id, table_name))
+        multi_predictions = self.type_multiclass_predictions.get((paper_id, table_name))
+        if prediction is not None:
+            multi_predictions = sorted(multi_predictions.items(), key=lambda x: x[1], reverse=True)
+            return prediction, [(k.name, v) for k, v in multi_predictions
+                                ]
+
 
 class LinkerEvaluator:
     def __init__(self, pipeline_logger, pc):
@@ -102,3 +167,18 @@ def on_taxonomy_topk(self, step, ext_id, topk):
 
     def top_matches(self, paper_id, table_name, row, col):
         return self.topk[(paper_id, table_name, row, col)]
+
+
+class FilteringEvaluator:
+    def __init__(self, pipeline_logger):
+        pipeline_logger.register("filtering::.*::filtered", self.on_filtered)
+        self.proposals = {}
+        self.which = {}
+        self.reason = pd.Series(dtype=str)
+
+    def on_filtered(self, step, proposals, which, reason, **kwargs):
+        _, filter_step, _ = step.split('::')
+        if filter_step != "compound_filtering":
+            self.proposals[filter_step] = pd.concat(self.proposals.get(filter_step, []) + [proposals])
+            self.which[filter_step] = pd.concat(self.which.get(filter_step, []) + [which])
+            self.reason = self.reason.append(reason)