mention_extractor.apply with clear=True fails if it's not the first run

CHANGELOG.rst

@@ -9,6 +9,8 @@ Fixed
 ^^^^^
 * `@kaikun213`_: Fix bug in table range difference calculations.
   (`#420 <https://github.com/HazyResearch/fonduer/pull/420>`_)
+* `@HiromuHota`_: mention_extractor.apply with clear=True now works even if it's not the first run.
+  (`#424 <https://github.com/HazyResearch/fonduer/pull/424>`_)
 
 0.8.2_ - 2020-04-28
 -------------------

src/fonduer/candidates/mentions.py

@@ -19,6 +19,7 @@
 from fonduer.candidates.models.temporary_context import TemporaryContext
 from fonduer.parser.models import Context, Document, Sentence
 from fonduer.utils.udf import UDF, UDFRunner
+from fonduer.utils.utils import get_dict_of_stable_id
 
 logger = logging.getLogger(__name__)
 
@@ -566,7 +567,7 @@ def apply(self, doc: Document, **kwargs: Any) -> Document:
         :param doc: A document to process.
         """
         # Get a dict of stable_id of contexts.
-        dict_of_stable_id: Dict[str, Context] = {}
+        dict_of_stable_id: Dict[str, Context] = get_dict_of_stable_id(doc)
 
         # Iterate over each mention class
         for i, mention_class in enumerate(self.mention_classes):

src/fonduer/utils/utils.py

@@ -2,7 +2,7 @@
 from builtins import range
 from typing import TYPE_CHECKING, Dict, Iterator, List, Set, Tuple, Type, Union
 
-from fonduer.parser.models import Document, Sentence
+from fonduer.parser.models import Context, Document, Sentence
 
 if TYPE_CHECKING:  # to prevent circular imports
     from fonduer.candidates.models import Candidate
@@ -64,3 +64,29 @@ def get_set_of_stable_ids(
             )
         )
     return set_of_stable_ids
+
+
+def get_dict_of_stable_id(doc: Document) -> Dict[str, Context]:
+    """Return a mapping of a stable_id to its context."""
+    return {
+        doc.stable_id: doc,
+        **{
+            c.stable_id: c
+            for a in [
+                "sentences",
+                "paragraphs",
+                "captions",
+                "cells",
+                "tables",
+                "sections",
+                "figures",
+            ]
+            for c in getattr(doc, a)
+        },
+        **{
+            c.stable_id: c
+            for s in doc.sentences
+            for a in ["spans", "implicit_spans"]
+            for c in getattr(s, a)
+        },
+    }

tests/e2e/test_incremental.py

@@ -92,6 +92,12 @@ def test_incremental():
     assert session.query(Part).count() == 11
     assert session.query(Temp).count() == 8
 
+    # Test if clear=True works
+    mention_extractor.apply(docs, parallelism=PARALLEL, clear=True)
+
+    assert session.query(Part).count() == 11
+    assert session.query(Temp).count() == 8
+
     # Candidate Extraction
     PartTemp = candidate_subclass("PartTemp", [Part, Temp])