Add PaperCollection

* add metadata when loading elastic documents from jsons
* load paper's text, tables and annotations

Author: Marcin Kardas

extract_texts.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import fire
 from sota_extractor2.data.elastic import Paper
 from pathlib import Path

sota_extractor2/data/elastic.py

@@ -12,6 +12,7 @@
 
 from sota_extractor2.data.doc_utils import get_text, content_in_section, group_content, set_ids_by_labels, read_html
 from .. import config
+from pathlib import Path
 
 
 def setup_default_connection():
@@ -108,6 +109,20 @@ class Fragment(Document):
     class Index:
         name = 'paper-fragments'
 
+    @classmethod
+    def from_json(cls, json):
+        if isinstance(json, str):
+            source = serializer.loads(json)
+        else:
+            source = json
+        data = dict(
+            _source = source,
+            _id = f"{source['paper_id']}_{source['order']}",
+            _index = 'paper-fragments',
+            _type = 'doc')
+        return cls.from_es(data)
+
+
     def __repr__(self):
         return f"# {self.header},\n" \
             f"{self.text}" \
@@ -125,7 +140,34 @@ class Index:
         name = 'papers'
 
     def to_json(self):
-        return serializer.dumps(self.to_dict())
+        data = self.to_dict()
+        return serializer.dumps(d)
+
+    @classmethod
+    def from_json(cls, json, paper_id=None):
+        if isinstance(json, str):
+            source = serializer.loads(json)
+        else:
+            source = json
+        fragments = source.pop('fragments', [])
+        data = dict(
+            _source = source,
+            _index = 'papers',
+            _type = 'doc')
+        if paper_id is not None:
+            data['_id'] = paper_id
+
+        paper = cls.from_es(data)
+        paper.fragments = Fragments([Fragment.from_json(f) for f in fragments])
+        return paper
+
+    @classmethod
+    def from_file(cls, path):
+        path = Path(path)
+        paper_id = path.stem
+        with open(path, "rt") as f:
+            json = f.read()
+        return cls.from_json(json, paper_id)
 
     def to_df(self):
         return pd.DataFrame({'header': [f.header for f in self.fragments],

sota_extractor2/data/json.py

@@ -69,11 +69,12 @@ def cut(s, length=20):
         vals = pprint.pformat({to_snake_case(k): cut(str(self[k]))  for k in self.keys()})
         return f"NodeWrap({vals})"
 
-def load_gql_dump(data_or_file):
+def load_gql_dump(data_or_file, compressed=True):
     if isinstance(data_or_file, dict):
         papers_data = data_or_file
     else:
-        with gzip.open(data_or_file, "rb") as f:
+        open_fn = gzip.open if compressed else open
+        with open_fn(data_or_file, "rt") as f:
                 papers_data = json.load(f)
     data = papers_data["data"]
     return {k:wrap_dict(v) for k,v in data.items()}

sota_extractor2/data/paper_collection.py

@@ -0,0 +1,76 @@
+from .elastic import Paper as PaperText
+from .table import Table, read_tables
+from .json import load_gql_dump
+from pathlib import Path
+import re
+
+class Paper:
+    def __init__(self, text, tables, annotations):
+        self.text = text
+        self.tables = tables
+        if annotations is not None:
+            self.gold_tags = annotations.gold_tags.strip()
+        else:
+            self.gold_tags = ''
+
+
+arxiv_version_re = re.compile(r"v\d+$")
+def clean_arxiv_version(arxiv_id):
+    return arxiv_version_re.sub("", arxiv_id)
+
+
+class PaperCollection:
+    def __init__(self, path, load_texts=True, load_tables=True):
+        self.path = path
+        self.load_texts = load_texts
+        self.load_tables = load_tables
+
+        if self.load_texts:
+            texts = self._load_texts()
+        else:
+            texts = {}
+
+        annotations = self._load_annotated_papers()
+        if self.load_tables:
+            tables = self._load_tables(annotations)
+        else:
+            tables = {}
+            annotations = {}
+        outer_join = set(texts).union(set(tables))
+
+        self._papers = {k: Paper(texts.get(k), tables.get(k), annotations.get(k)) for k in outer_join}
+
+    def __len__(self):
+        return len(self._papers)
+
+    def __getitem__(self, idx):
+        return self._papers[idx]
+
+    def __iter__(self):
+        return iter(self._papers)
+
+    def _load_texts(self):
+        texts = {}
+
+        for f in (self.path / "texts").glob("**/*.json"):
+            text = PaperText.from_file(f)
+            texts[clean_arxiv_version(text.meta.id)] = text
+        return texts
+
+
+    def _load_tables(self, annotations):
+        tables = {}
+
+        for f in (self.path / "tables").glob("**/metadata.json"):
+            paper_dir = f.parent
+            tbls = read_tables(paper_dir, annotations)
+            tables[clean_arxiv_version(paper_dir.name)] = tbls
+        return tables
+
+    def _load_annotated_papers(self):
+        dump = load_gql_dump(self.path / "structure-annotations.json.gz", compressed=True)["allPapers"]
+        annotations = {}
+        for a in dump:
+            arxiv_id = clean_arxiv_version(a.arxiv_id)
+            annotations[arxiv_id] = a
+        return annotations

sota_extractor2/data/table.py

@@ -0,0 +1,46 @@
+import pandas as pd
+import json
+from pathlib import Path
+from dataclasses import dataclass
+from typing import List
+from ..helpers.jupyter import display_table
+
+@dataclass
+class Cell:
+    value: str
+    gold_tags: str = ''
+    refs: List[str] = None
+
+
+class Table:
+    def __init__(self, df, caption=None, figure_id=None, annotations=None):
+        self.df = df
+        self.caption = caption
+        self.figure_id = figure_id
+        self.df = df.applymap(lambda x: Cell(value=x))
+        if annotations is not None:
+            self.gold_tags = annotations.gold_tags.strip()
+            rows, cols = annotations.matrix_gold_tags.shape
+            for r in range(rows):
+                for c in range(cols):
+                    self.df.iloc[r,c].gold_tags = annotations.matrix_gold_tags.iloc[r,c].strip()
+        else:
+            self.gold_tags = ''
+
+    @classmethod
+    def from_file(cls, path, metadata, annotations=None):
+        try:
+            df = pd.read_csv(path, header=None, dtype=str).fillna('')
+        except pd.errors.EmptyDataError:
+            df = pd.DataFrame()
+        return cls(df, metadata.get('caption'), metadata.get('figure_id'), annotations)
+
+    def display(self):
+
+        display_table(self.df.applymap(lambda x: x.value).values, self.df.applymap(lambda x: x.gold_tags).values)
+
+def read_tables(path, annotations):
+    path = Path(path)
+    with open(path / "metadata.json", "r") as f:
+        metadata = json.load(f)
+    return [Table.from_file(path / m["filename"], m, annotations.get(path.name)) for m in metadata]

sota_extractor2/helpers/jupyter.py

@@ -16,7 +16,10 @@ def display_table(table, structure=None):
         matrix - 2d ndarray with cell values
         strucutre - 2d ndarray with structure annotation
     """
-    matrix = table.matrix
+    if hasattr(table, 'matrix'):
+        matrix = table.matrix
+    else:
+        matrix = table
     if structure is None: structure = table.matrix_gold_tags
     html = []
     html.append('<link href="http://10.0.1.145:8001/static/css/main.bd3d2d63.chunk.css" rel="stylesheet">')
@@ -29,4 +32,4 @@ def display_table(table, structure=None):
         html.append("</tr>")
     html.append("</table>")
     html.append('</div>')
-    display_html("\n".join(html))
+    display_html("\n".join(html))