Merge pull request #39 from allenai/soldni/more_eval

Fixes for grobid eval

papermage/parsers/grobid_parser.py

@@ -1,21 +1,21 @@
 """
 
-@geli-gel
+@geli-gel, @amanpreet692, @soldni
 
 """
 import json
 import os
 import re
+import warnings
 import xml.etree.ElementTree as et
 from collections import defaultdict
 from tempfile import NamedTemporaryFile
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Dict, List, Optional, Tuple, cast
 
 import numpy as np
 from grobid_client.grobid_client import GrobidClient
 
 from papermage.magelib import (
-    Annotation,
     Document,
     Entity,
     Metadata,
@@ -67,9 +67,9 @@
 }
 
 
-def find_zero_spans(array):
-    # Add a sentinel value at the end to not miss end of last span
-    array = np.append(array, 1)
+def find_contiguous_ones(array):
+    # Add a sentinel value at the beginning/end
+    array = np.concatenate([[0], array, [0]])
 
     # Find the indexes where the array changes
     diff_indices = np.where(np.diff(array) != 0)[0] + 1
@@ -80,7 +80,7 @@ def find_zero_spans(array):
     spans = list(zip(zero_start_indices, zero_end_indices))
 
     # Exclude the spans with no element
-    spans = [(start, end + 1) for start, end in spans if end - start >= 0]
+    spans = [(start - 1, end) for start, end in spans if end - start >= 0]
 
     return spans
 
@@ -95,13 +95,14 @@ class GrobidFullParser(Parser):
     > docker run -t --rm -p 8070:8070 lfoppiano/grobid:0.7.2
     """
 
-    def __init__(self, grobid_config: Optional[dict] = None, check_server: bool = True):
-        self.grobid_config = grobid_config or {
+    def __init__(self, check_server: bool = True, **grobid_config: Any):
+        self.grobid_config = {
             "grobid_server": "http://localhost:8070",
             "batch_size": 1000,
             "sleep_time": 5,
             "timeout": 60,
             "coordinates": sorted(set((*GROBID_VILA_MAP.keys(), "s", "ref", "body", "item", "persName"))),
+            **grobid_config,
         }
         assert "coordinates" in self.grobid_config, "Grobid config must contain 'coordinates' key"
 
@@ -113,14 +114,19 @@ def __init__(self, grobid_config: Optional[dict] = None, check_server: bool = Tr
 
         os.remove(config_path)
 
-    def parse(self, input_pdf_path: str, doc: Document, xml_out_dir: Optional[str] = None) -> Document:
+    def parse(  # type: ignore
+        self,
+        input_pdf_path: str,
+        doc: Document,
+        xml_out_dir: Optional[str] = None
+    ) -> Document:
         assert doc.symbols != ""
         for field in REQUIRED_DOCUMENT_FIELDS:
             assert field in doc.fields
 
         (_, _, xml) = self.client.process_pdf(
-            "processFulltextDocument",
-            input_pdf_path,
+            service="processFulltextDocument",
+            pdf_file=input_pdf_path,
             generateIDs=False,
             consolidate_header=False,
             consolidate_citations=False,
@@ -137,16 +143,16 @@ def parse(self, input_pdf_path: str, doc: Document, xml_out_dir: Optional[str] =
             with open(xmlfile, "w") as f_out:
                 f_out.write(xml)
 
-        self._parse_xml_onto_doc(xml, doc)
+        self._parse_xml_onto_doc(xml=xml, doc=doc)
 
-        for p in doc.p:
+        for p in getattr(doc, "p", []):
             grobid_text_elems = [s.metadata["grobid_text"] for s in p.s]
             grobid_text = " ".join(filter(lambda text: isinstance(text, str), grobid_text_elems))
             p.metadata["grobid_text"] = grobid_text
 
         # add vila-like entities
-        doc.annotate_entity(entities=self._make_vila_groups(doc), field_name="vila_entities")
-
+        vila_entities = self._make_vila_groups(doc)
+        doc.annotate_entity(entities=vila_entities, field_name="vila_entities")
         return doc
 
     def _make_spans_from_boxes(self, doc: Document, entity: Entity) -> List[Span]:
@@ -173,18 +179,20 @@ def _make_entities_of_type(
         ]
         return entities
 
-    def _update_reserved_positions(self, reserved_positions: np.ndarray, entities: List[Entity]) -> List[Entity]:
+    def _update_reserved_positions(
+        self, reserved_positions: np.ndarray, entities: List[Entity]
+    ) -> Tuple[List[Entity], np.ndarray]:
         new_entities: List[Entity] = []
         for ent in entities:
             new_spans = []
             for span in ent.spans:
                 already_reserved = reserved_positions[span.start : span.end]
-                for start, end in find_zero_spans(already_reserved):
-                    new_spans.append(Span(start=start, end=end))
+                for start, end in find_contiguous_ones(~already_reserved):
+                    new_spans.append(Span(start=start + span.start, end=end + span.start))
                 reserved_positions[span.start : span.end] = True
             if new_spans:
                 new_entities.append(Entity(spans=new_spans, boxes=ent.boxes, metadata=ent.metadata))
-        return new_entities
+        return new_entities, reserved_positions
 
     def _make_vila_groups(self, doc: Document) -> List[Entity]:
         ents: List[Entity] = []
@@ -196,106 +204,104 @@ def _make_vila_groups(self, doc: Document) -> List[Entity]:
 
         if h := getattr(doc, "author", []):
             h_ = self._make_entities_of_type(doc=doc, entities=h, entity_type="Author", id_offset=len(ents))
-            h_ = self._update_reserved_positions(reserved_positions, h_)
-            ents.extend(h_)
+            h__, reserved_positions = self._update_reserved_positions(reserved_positions, h_)
+            ents.extend(h__)
 
         if a := getattr(doc, "abstract", []):
             a_ = self._make_entities_of_type(doc=doc, entities=a, entity_type="Abstract", id_offset=len(ents))
-            a_ = self._update_reserved_positions(reserved_positions, a_)
-            ents.extend(a_)
+            a__, reserved_positions = self._update_reserved_positions(reserved_positions, a_)
+            ents.extend(a__)
 
         if _ := getattr(doc, "keywords", []):
             # keywords has no coordinates, so we can't recover their positions!
             pass
 
         if s := getattr(doc, "head", []):
             s_ = self._make_entities_of_type(doc=doc, entities=s, entity_type="Section", id_offset=len(ents))
-            s_ = self._update_reserved_positions(reserved_positions, s_)
-            ents.extend(s_)
+            s__, reserved_positions = self._update_reserved_positions(reserved_positions, s_)
+            ents.extend(s__)
 
-        if l := getattr(doc, "list", []):
-            l_ = self._make_entities_of_type(doc=doc, entities=l, entity_type="List", id_offset=len(ents))
-            l_ = self._update_reserved_positions(reserved_positions, l_)
-            ents.extend(l_)
+        if t := getattr(doc, "list", []):
+            t_ = self._make_entities_of_type(doc=doc, entities=t, entity_type="List", id_offset=len(ents))
+            t__, reserved_positions = self._update_reserved_positions(reserved_positions, t_)
+            ents.extend(t__)
 
         if b := getattr(doc, "biblStruct", []):
             b_ = self._make_entities_of_type(doc=doc, entities=b, entity_type="Bibliography", id_offset=len(ents))
-            b_ = self._update_reserved_positions(reserved_positions, b_)
-            ents.extend(b_)
+            b__, reserved_positions = self._update_reserved_positions(reserved_positions, b_)
+            ents.extend(b__)
 
         if e := getattr(doc, "formula", []):
             e_ = self._make_entities_of_type(doc=doc, entities=e, entity_type="Equation", id_offset=len(ents))
-            e_ = self._update_reserved_positions(reserved_positions, e_)
-            ents.extend(e_)
+            e__, reserved_positions = self._update_reserved_positions(reserved_positions, e_)
+            ents.extend(e__)
 
-        if figures := getattr(doc, "figure", []):
-            for figure in figures:
-                current_boxes = [Box(l=b.l, t=b.t, w=b.w, h=b.h, page=b.page) for b in figure.boxes]
+        if figs := getattr(doc, "figure", []):
+            for fig in figs:
+                current_boxes = [Box(l=b.l, t=b.t, w=b.w, h=b.h, page=b.page) for b in fig.boxes]
 
                 if "figDesc" in doc.fields:
-                    caption_boxes = [b for d in doc.find_by_box(figure, "figDesc") for b in d.boxes]
+                    caption_boxes = [b for d in doc.find_by_box(fig, "figDesc") for b in d.boxes]
                     current_boxes = [b for b in current_boxes if b not in caption_boxes]
 
                 if "table" in doc.fields:
-                    table_boxes = [b for d in doc.find_by_box(figure, "table") for b in d.boxes]
+                    table_boxes = [b for d in doc.find_by_box(fig, "table") for b in d.boxes]
                     current_boxes = [b for b in current_boxes if b not in table_boxes]
 
                 if not current_boxes:
                     continue
 
-                new_figure = Entity(
-                    spans=None,
+                new_fig = Entity(
+                    spans=self._make_spans_from_boxes(doc, Entity(boxes=current_boxes)),
                     boxes=current_boxes,
-                    metadata=Metadata(**figure.metadata.to_json(), type="Figure", id=len(ents)),
+                    metadata=Metadata(**fig.metadata.to_json(), label="Figure", id=len(ents)),
                 )
-                ents.append(new_figure)
+                new_figs, reserved_positions = self._update_reserved_positions(reserved_positions, [new_fig])
+                ents.extend(new_figs)
 
         if t := getattr(doc, "table", []):
             t_ = self._make_entities_of_type(doc=doc, entities=t, entity_type="Table", id_offset=len(ents))
-            t_ = self._update_reserved_positions(reserved_positions, t_)
-            ents.extend(t_)
+            t__, reserved_positions = self._update_reserved_positions(reserved_positions, t_)
+            ents.extend(t__)
 
         if c := getattr(doc, "figDesc", []):
             c_ = self._make_entities_of_type(doc=doc, entities=c, entity_type="Caption", id_offset=len(ents))
-            c_ = self._update_reserved_positions(reserved_positions, c_)
-            ents.extend(c_)
+            c__, reserved_positions = self._update_reserved_positions(reserved_positions, c_)
+            ents.extend(c__)
 
         if _ := getattr(doc, "note", []):
             # notes have no coordinates, so we can't recover their positions!
             pass
 
         if p := getattr(doc, "p", []):
             p_ = self._make_entities_of_type(doc=doc, entities=p, entity_type="Paragraph", id_offset=len(ents))
-            p_ = self._update_reserved_positions(reserved_positions, p_)
-            ents.extend(p_)
+            p__, reserved_positions = self._update_reserved_positions(reserved_positions, p_)
+            ents.extend(p__)
 
         return ents
 
     def _parse_xml_onto_doc(self, xml: str, doc: Document) -> Document:
-        xml_root = et.fromstring(xml)
+        try:
+            xml_root = et.fromstring(xml)
+        except Exception as e:
+            if xml == "[GENERAL] An exception occurred while running Grobid.":
+                warnings.warn("Grobid returned an error; check server logs")
+                return doc
+            raise e
 
         all_box_groups = self._get_box_groups(xml_root)
         for field, box_groups in all_box_groups.items():
-            # span_groups = box_groups_to_span_groups(
-            #     box_groups=box_groups, doc=doc, center=True
-            # )
-            # assert len(box_groups) == len(span_groups), (
-            #     f"Annotations and SpanGroups for {field} are not the same length"
-            # )
-            # for bg, sg in zip(box_groups, span_groups):
-            #     sg.metadata = bg.metadata
-            #
-            # # note for if/when adding in relations between mention sources and
-            # # bib targets: big_entries metadata contains original grobid id
-            # # attached to the Annotation.
             doc.annotate_entity(field_name=field, entities=box_groups)
 
         return doc
 
-    def _xml_coords_to_boxes(self, coords_attribute: str, page_sizes: dict):
+    def _xml_coords_to_boxes(self, coords_attribute: str, page_sizes: dict) -> List[Box]:
         coords_list = coords_attribute.split(";")
         boxes = []
         for coords in coords_list:
+            if coords == "":
+                # this page has no coordinates
+                continue
             pg, x, y, w, h = coords.split(",")
             proper_page = int(pg) - 1
             boxes.append(
@@ -305,7 +311,7 @@ def _xml_coords_to_boxes(self, coords_attribute: str, page_sizes: dict):
             )
         return boxes
 
-    def _get_box_groups(self, root: et.Element) -> Dict[str, List[Annotation]]:
+    def _get_box_groups(self, root: et.Element) -> Dict[str, List[Entity]]:
         page_size_root = root.find(".//tei:facsimile", NS)
         assert page_size_root is not None, "No facsimile found in Grobid XML"
 
@@ -314,7 +320,7 @@ def _get_box_groups(self, root: et.Element) -> Dict[str, List[Annotation]]:
         for data in page_size_data:
             page_sizes[int(data.attrib["n"]) - 1] = [float(data.attrib["lrx"]), float(data.attrib["lry"])]
 
-        all_boxes: Dict[str, List[Annotation]] = defaultdict(list)
+        all_boxes: Dict[str, List[Entity]] = defaultdict(list)
 
         for field in self.grobid_config["coordinates"]:
             structs = root.findall(f".//tei:{field}", NS)
@@ -326,7 +332,11 @@ def _get_box_groups(self, root: et.Element) -> Dict[str, List[Annotation]]:
                 if coords_str == "":
                     continue
 
-                boxes = self._xml_coords_to_boxes(coords_str, page_sizes)
+                if not (boxes := self._xml_coords_to_boxes(coords_str, page_sizes)):
+                    # we check if the boxes are empty because sometimes users monkey-patch
+                    # _xml_coords_to_boxes to filter by page
+                    continue
+
                 metadata_dict: Dict[str, Any] = {
                     f"grobid_{re.sub(r'[^a-zA-Z0-9_]+', '_', k)}": v
                     for k, v in struct.attrib.items()
@@ -336,7 +346,6 @@ def _get_box_groups(self, root: et.Element) -> Dict[str, List[Annotation]]:
                 metadata_dict["grobid_text"] = struct.text
                 metadata = Metadata.from_json(metadata_dict)
                 box_group = Entity(boxes=boxes, metadata=metadata)
-                print(box_group)
                 all_boxes[field].append(box_group)
 
         return all_boxes
@@ -354,6 +363,6 @@ def _get_box_groups(self, root: et.Element) -> Dict[str, List[Annotation]]:
     doc = PDFPlumberParser().parse(opts.pdf_path)
     doc = GrobidFullParser().parse(opts.pdf_path, doc)
 
-    for p in doc.p:
+    for p in getattr(doc, "p", []):
         for s in p.s:
             print(s.metadata.grobid_text)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = 'papermage'
-version = '0.10.0'
+version = '0.11.0'
 description = 'Papermage. Casting magic over scientific PDFs.'
 license = {text = 'Apache-2.0'}
 readme = 'README.md'