Use new annotations format

Author: Marcin Kardas

axcell/data/json.py

@@ -71,13 +71,62 @@ 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 _annotations_to_gql(annotations):
+    nodes = []
+    for a in annotations:
+        tables = []
+        for t in a['tables']:
+            tags = []
+            if t['leaderboard']:
+                tags.append('leaderboard')
+            if t['ablation']:
+                tags.append('ablation')
+            if not tags:
+                tags = ['irrelevant']
+
+            records = {}
+            for r in t['records']:
+                d = dict(r)
+                del d['row']
+                del d['column']
+                records[f'{r["row"]}.{r["column"]}'] = d
+            table = {
+                'node': {
+                    'name': f'table_{t["index"] + 1:02}.csv',
+                    'datasetText': t['dataset_text'],
+                    'notes': '',
+                    'goldTags': ' '.join(tags),
+                    'matrixGoldTags': t['segmentation'],
+                    'cellsSotaRecords': json.dumps(records),
+                    'parser': 'latexml'
+                }
+            }
+            tables.append(table)
+        node = {
+            'arxivId': a['arxiv_id'],
+            'goldTags': a['fold'],
+            'tableSet': {'edges': tables}
+        }
+        nodes.append({'node': node})
+    return {
+        'data': {
+            'allPapers': {
+                'edges': nodes
+            }
+        }
+    }
+
+
 def load_gql_dump(data_or_file, compressed=True):
-    if isinstance(data_or_file, dict):
+    if isinstance(data_or_file, dict) or isinstance(data_or_file, list):
         papers_data = data_or_file
     else:
         open_fn = gzip.open if compressed else open
         with open_fn(data_or_file, "rt") as f:
-                papers_data = json.load(f)
+            papers_data = json.load(f)
+    if "data" not in papers_data:
+        papers_data = _annotations_to_gql(papers_data)
     data = papers_data["data"]
     return {k:wrap_dict(v) for k,v in data.items()}
 

axcell/data/paper_collection.py

@@ -75,23 +75,32 @@ def _load_tables(path, annotations, jobs, migrate):
     return {f.parent.name: tbls for f, tbls in zip(files, tables)}
 
 
+def _gql_dump_to_annotations(dump):
+    annotations = {remove_arxiv_version(a.arxiv_id): a for a in dump}
+    annotations.update({a.arxiv_id: a for a in dump})
+    return annotations
+
 def _load_annotated_papers(data_or_path):
-    if isinstance(data_or_path, dict):
+    if isinstance(data_or_path, dict) or isinstance(data_or_path, list):
         compressed = False
     else:
         compressed = data_or_path.suffix == ".gz"
     dump = load_gql_dump(data_or_path, compressed=compressed)["allPapers"]
-    annotations = {remove_arxiv_version(a.arxiv_id): a for a in dump}
-    annotations.update({a.arxiv_id: a for a in dump})
-    return annotations
+    return _gql_dump_to_annotations(dump)
 
 
 class PaperCollection(UserList):
     def __init__(self, data=None):
         super().__init__(data)
 
     @classmethod
-    def from_files(cls, path, annotations_path=None, load_texts=True, load_tables=True, load_annotations=True, jobs=-1, migrate=False):
+    def from_files(cls, path, annotations=None, load_texts=True, load_tables=True, jobs=-1):
+        return cls._from_files(path, annotations=annotations, annotations_path=None,
+                               load_texts=load_texts, load_tables=load_tables, load_annotations=False,
+                               jobs=jobs)
+
+    @classmethod
+    def _from_files(cls, path, annotations=None, annotations_path=None, load_texts=True, load_tables=True, load_annotations=True, jobs=-1, migrate=False):
         path = Path(path)
         if annotations_path is None:
             annotations_path = path / "structure-annotations.json"
@@ -102,7 +111,10 @@ def from_files(cls, path, annotations_path=None, load_texts=True, load_tables=Tr
         else:
             texts = {}
 
-        annotations = {}
+        if annotations is None:
+            annotations = {}
+        else:
+            annotations = _load_annotated_papers(annotations)
         if load_tables:
             if load_annotations:
                 annotations = _load_annotated_papers(annotations_path)

extract_tables.py

@@ -348,7 +348,7 @@ def remove_footnotes(soup):
 
 
 def extract_tables(html):
-    soup = BeautifulSoup(html, "lxml", from_encoding="utf-8")
+    soup = BeautifulSoup(html, "lxml")
     set_ids_by_labels(soup)
     fix_span_tables(soup)
     fix_th(soup)