Merge branch 'master' of https://github.com/paperswithcode/paper-extractor

Author: PiotrCzapla

Makefile

@@ -72,8 +72,12 @@ $(ANNOTATIONS_DIR)/evaluation-tables.json.gz:
 	$(shell mkdir -p "$(ANNOTATIONS_DIR)")
 	wget https://paperswithcode.com/media/about/evaluation-tables.json.gz -O $@
 
+.PHONY: pull_images
+pull_images:
+	docker pull arxivvanity/engrafo:b3db888fefa118eacf4f13566204b68ce100b3a6
+	docker pull zenika/alpine-chrome:73
 
-.PHONY : clean
+.PHONY: clean
 clean :
 	cd "$(ANNOTATIONS_DIR)" && rm -f *.json *.csv
 	#rm -f *.gz

README.md

@@ -20,8 +20,9 @@ Directory structure:
 ```
 
 
-To preprocess data and extract tables, run:
+To preprocess data and extract tables and texts, run:
 ```
+make pull_images
 conda env create -f environment.yml
 source activate xtables
 make -j 8 -i extract_all > stdout.log 2> stderr.log

helpers.py

@@ -0,0 +1,43 @@
+from fire import Fire
+from pathlib import Path
+from sota_extractor2.data.paper_collection import PaperCollection
+from sota_extractor2.data.structure import CellEvidenceExtractor
+from elasticsearch_dsl import connections
+from tqdm import tqdm
+import pandas as pd
+from joblib import delayed, Parallel
+
+class Helper:
+    def split_pc_pickle(self, path, outdir="pc-parts", parts=8):
+        outdir = Path(outdir)
+        outdir.mkdir(parents=True, exist_ok=True)
+        pc = PaperCollection.from_pickle(path)
+        step = (len(pc) + parts - 1) // parts
+        for idx, i in enumerate(range(0, len(pc), step)):
+            part = PaperCollection(pc[i:i + step])
+            part.to_pickle(outdir / f"pc-part-{idx:02}.pkl")
+
+    def _evidences_for_pc(self, path):
+        path = Path(path)
+        pc = PaperCollection.from_pickle(path)
+        cell_evidences = CellEvidenceExtractor()
+        connections.create_connection(hosts=['10.0.1.145'], timeout=20)
+        raw_evidences = []
+        for paper in tqdm(pc):
+            raw_evidences.append(cell_evidences(paper, paper.tables, paper_limit=100, corpus_limit=20))
+        raw_evidences = pd.concat(raw_evidences)
+        path = path.with_suffix(".evidences.pkl")
+        raw_evidences.to_pickle(path)
+
+    def evidences_for_pc(self, pattern="pc-parts/pc-part-??.pkl", jobs=-1):
+        pickles = sorted(Path(".").glob(pattern))
+        Parallel(backend="multiprocessing", n_jobs=jobs)(delayed(self._evidences_for_pc)(path) for path in pickles)
+
+    def merge_evidences(self, output="evidences.pkl", pattern="pc-parts/pc-part-*.evidences.pkl"):
+        pickles = sorted(Path(".").glob(pattern))
+        evidences = [pd.read_pickle(pickle) for pickle in pickles]
+        evidences = pd.concat(evidences)
+        evidences.to_pickle(output)
+
+
+if __name__ == "__main__": Fire(Helper())

sota_extractor2/data/table.py

@@ -146,6 +146,10 @@ def set_tags(self, tags):
                 # todo: change gold_tags to tags to avoid confusion
                 self.df.iloc[r,c].gold_tags = cell.strip()
 
+    @property
+    def shape(self):
+        return self.df.shape
+
     @property
     def matrix(self):
         return self.df.applymap(lambda x: x.value)

sota_extractor2/helpers/training.py

@@ -1,10 +1,14 @@
 
-def set_seed(seed, name, quiet=False):
+def set_seed(seed, name, quiet=False, all_gpus=True):
     import torch
     import numpy as np
+    import random
     if not quiet:
         print(f"Setting {name} seed to {seed}")
     torch.manual_seed(seed)
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
-    np.random.seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    if all_gpus:
+        torch.cuda.manual_seed_all(seed)

sota_extractor2/models/linking/context_search.py

@@ -138,13 +138,17 @@
     'LibriSpeech dev-other': ['libri speech dev other', 'libri speech', 'dev', 'other', 'dev other', 'development', 'noisy'],
 })
 
+tasks = {}
+
 # escaped_ws_re = re.compile(r'\\\s+')
 # def name_to_re(name):
 #     return re.compile(r'(?:^|\s+)' + escaped_ws_re.sub(r'\\s*', re.escape(name.strip())) + r'(?:$|\s+)', re.I)
 
 #all_datasets = set(k for k,v in merged_p.items() if k != '' and not re.match("^\d+$", k) and v.get('NOMATCH', 0.0) < 0.9)
 all_datasets = set(normalize_cell_ws(normalize_dataset(y)) for x in datasets.values() for y in x)
 all_metrics = set(normalize_cell_ws(y) for x in metrics.values() for y in x)
+all_tasks = set(normalize_cell_ws(normalize_dataset(y)) for x in tasks.values() for y in x)
+
 #all_metrics = set(metrics_p.keys())
 
 # all_datasets_re = {x:name_to_re(x) for x in all_datasets}
@@ -187,6 +191,7 @@ def find_names(text, names_trie):
 
 all_datasets_trie = make_trie(all_datasets)
 all_metrics_trie = make_trie(all_metrics)
+all_tasks_trie = make_trie(all_tasks)
 
 
 def find_datasets(text):
@@ -195,18 +200,23 @@ def find_datasets(text):
 def find_metrics(text):
     return find_names(text, all_metrics_trie)
 
+def find_tasks(text):
+    return find_names(text, all_tasks_trie)
+
 def dummy_item(reason):
     return pd.DataFrame(dict(dataset=[reason], task=[reason], metric=[reason], evidence=[""], confidence=[0.0]))
 
 
 
 @njit
-def compute_logprobs(taxonomy, reverse_merged_p, reverse_metrics_p, dss, mss, noise, ms_noise, ds_pb, ms_pb, logprobs):
+def compute_logprobs(taxonomy, reverse_merged_p, reverse_metrics_p, reverse_task_p,
+                     dss, mss, tss, noise, ms_noise, ts_noise, ds_pb, ms_pb, ts_pb, logprobs):
     empty = typed.Dict.empty(types.unicode_type, types.float64)
     for i, (task, dataset, metric) in enumerate(taxonomy):
         logprob = 0.0
         short_probs = reverse_merged_p.get(dataset, empty)
         met_probs = reverse_metrics_p.get(metric, empty)
+        task_probs = reverse_task_p.get(task, empty)
         for ds in dss:
             #                 for abbrv, long_form in abbrvs.items():
             #                     if ds == abbrv:
@@ -216,17 +226,21 @@ def compute_logprobs(taxonomy, reverse_merged_p, reverse_metrics_p, dss, mss, no
             logprob += np.log(noise * ds_pb + (1 - noise) * short_probs.get(ds, 0.0))
         for ms in mss:
             logprob += np.log(ms_noise * ms_pb + (1 - ms_noise) * met_probs.get(ms, 0.0))
+        for ts in tss:
+            logprob += np.log(ts_noise * ts_pb + (1 - ts_noise) * task_probs.get(ts, 0.0))
         logprobs[i] += logprob
         #logprobs[(dataset, metric)] = logprob
 
 
 class ContextSearch:
-    def __init__(self, taxonomy, context_noise=(0.5, 0.2, 0.1), metrics_noise=None, ds_pb=0.001, ms_pb=0.01, debug_gold_df=None):
+    def __init__(self, taxonomy, context_noise=(0.5, 0.2, 0.1), metrics_noise=None, task_noise=None,
+                 ds_pb=0.001, ms_pb=0.01, ts_pb=0.01, debug_gold_df=None):
         merged_p = \
         get_probs({k: Counter([normalize_cell(normalize_dataset(x)) for x in v]) for k, v in datasets.items()})[1]
         metrics_p = \
         get_probs({k: Counter([normalize_cell(normalize_dataset(x)) for x in v]) for k, v in metrics.items()})[1]
-
+        tasks_p = \
+        get_probs({k: Counter([normalize_cell(normalize_dataset(x)) for x in v]) for k, v in tasks.items()})[1]
 
         self.queries = {}
         self.taxonomy = taxonomy
@@ -236,10 +250,13 @@ def __init__(self, taxonomy, context_noise=(0.5, 0.2, 0.1), metrics_noise=None,
         self.extract_acronyms = AcronymExtractor()
         self.context_noise = context_noise
         self.metrics_noise = metrics_noise if metrics_noise else context_noise
+        self.task_noise = task_noise if task_noise else context_noise
         self.ds_pb = ds_pb
         self.ms_pb = ms_pb
+        self.ts_pb = ts_pb
         self.reverse_merged_p = self._numba_update_nested_dict(reverse_probs(merged_p))
         self.reverse_metrics_p = self._numba_update_nested_dict(reverse_probs(metrics_p))
+        self.reverse_tasks_p = self._numba_update_nested_dict(reverse_probs(tasks_p))
         self.debug_gold_df = debug_gold_df
 
     def _numba_update_nested_dict(self, nested):
@@ -256,29 +273,33 @@ def _numba_extend_list(self, lst):
             l.append(x)
         return l
 
-    def compute_context_logprobs(self, context, noise, ms_noise, logprobs):
+    def compute_context_logprobs(self, context, noise, ms_noise, ts_noise, logprobs):
         context = context or ""
         abbrvs = self.extract_acronyms(context)
         context = normalize_cell_ws(normalize_dataset(context))
         dss = set(find_datasets(context)) | set(abbrvs.keys())
         mss = set(find_metrics(context))
+        tss = set(find_tasks(context))
         dss -= mss
+        dss -= tss
         dss = [normalize_cell(ds) for ds in dss]
         mss = [normalize_cell(ms) for ms in mss]
+        tss = [normalize_cell(ts) for ts in tss]
         ###print("dss", dss)
         ###print("mss", mss)
         dss = self._numba_extend_list(dss)
         mss = self._numba_extend_list(mss)
-        compute_logprobs(self._taxonomy, self.reverse_merged_p, self.reverse_metrics_p,
-                         dss, mss, noise, ms_noise, self.ds_pb, self.ms_pb, logprobs)
+        tss = self._numba_extend_list(tss)
+        compute_logprobs(self._taxonomy, self.reverse_merged_p, self.reverse_metrics_p, self.reverse_tasks_p,
+                         dss, mss, tss, noise, ms_noise, ts_noise, self.ds_pb, self.ms_pb, self.ts_pb, logprobs)
 
     def match(self, contexts):
         assert len(contexts) == len(self.context_noise)
         n = len(self._taxonomy)
         context_logprobs = np.zeros(n)
 
-        for context, noise, ms_noise in zip(contexts, self.context_noise, self.metrics_noise):
-            self.compute_context_logprobs(context, noise, ms_noise, context_logprobs)
+        for context, noise, ms_noise, ts_noise in zip(contexts, self.context_noise, self.metrics_noise, self.task_noise):
+            self.compute_context_logprobs(context, noise, ms_noise, ts_noise, context_logprobs)
         keys = self.taxonomy.taxonomy
         logprobs = context_logprobs
         #keys, logprobs = zip(*context_logprobs.items())

sota_extractor2/models/structure/experiment.py

@@ -20,15 +20,40 @@ class Labels(Enum):
     EMPTY=5
 
 
+class LabelsExt(Enum):
+    OTHER=0
+    PARAMS=6
+    TASK=7
+    DATASET=1
+    SUBDATASET=8
+    PAPER_MODEL=2
+    BEST_MODEL=9
+    ENSEMBLE_MODEL=10
+    COMPETING_MODEL=3
+    METRIC=4
+    EMPTY=5
+
+
 label_map = {
     "dataset": Labels.DATASET.value,
     "dataset-sub": Labels.DATASET.value,
     "model-paper": Labels.PAPER_MODEL.value,
     "model-best": Labels.PAPER_MODEL.value,
     "model-ensemble": Labels.PAPER_MODEL.value,
     "model-competing": Labels.COMPETING_MODEL.value,
-    "dataset-metric": Labels.METRIC.value,
-#    "model-params": Labels.PARAMS.value
+    "dataset-metric": Labels.METRIC.value
+}
+
+label_map_ext = {
+    "dataset": LabelsExt.DATASET.value,
+    "dataset-sub": LabelsExt.SUBDATASET.value,
+    "model-paper": LabelsExt.PAPER_MODEL.value,
+    "model-best": LabelsExt.BEST_MODEL.value,
+    "model-ensemble": LabelsExt.ENSEMBLE_MODEL.value,
+    "model-competing": LabelsExt.COMPETING_MODEL.value,
+    "dataset-metric": LabelsExt.METRIC.value,
+    "model-params": LabelsExt.PARAMS.value,
+    "dataset-task": LabelsExt.TASK.value
 }
 
 # put here to avoid recompiling, used only in _limit_context
@@ -63,6 +88,7 @@ class Experiment:
     remove_num: bool = True
     drop_duplicates: bool = True
     mark_this_paper: bool = False
+    distinguish_model_source: bool = True
 
     results: dict = dataclasses.field(default_factory=dict)
 
@@ -219,6 +245,8 @@ def _transform_df(self, df):
             df = df.replace(re.compile(r"(^|[ ])\d+(\b|%)"), " xxnum ")
         df = df.replace(re.compile(r"\bdata set\b"), " dataset ")
         df["label"] = df["cell_type"].apply(lambda x: label_map.get(x, 0))
+        if not self.distinguish_model_source:
+            df["label"] = df["label"].apply(lambda x: x if x != Labels.COMPETING_MODEL.value else Labels.PAPER_MODEL.value)
         df["label"] = pd.Categorical(df["label"])
         return df
 
@@ -228,13 +256,15 @@ def transform_df(self, *dfs):
             return transformed[0]
         return transformed
 
-    def _set_results(self, prefix, preds, true_y):
+    def _set_results(self, prefix, preds, true_y, true_y_ext=None):
         m = metrics(preds, true_y)
         r = {}
         r[f"{prefix}_accuracy"] = m["accuracy"]
         r[f"{prefix}_precision"] = m["precision"]
         r[f"{prefix}_recall"] = m["recall"]
         r[f"{prefix}_cm"] = confusion_matrix(true_y, preds, labels=[x.value for x in Labels]).tolist()
+        if true_y_ext is not None:
+            r[f"{prefix}_cm_full"] = confusion_matrix(true_y_ext, preds, labels=[x.value for x in LabelsExt]).tolist()
         self.update_results(**r)
 
     def evaluate(self, model, train_df, valid_df, test_df):
@@ -253,17 +283,19 @@ def evaluate(self, model, train_df, valid_df, test_df):
                 true_y = vote_results["true"]
             else:
                 true_y = tdf["label"]
-            self._set_results(prefix, preds, true_y)
+                true_y_ext = tdf["cell_type"].apply(lambda x: label_map_ext.get(x, 0))
+            self._set_results(prefix, preds, true_y, true_y_ext)
 
-    def show_results(self, *ds, normalize=True):
+    def show_results(self, *ds, normalize=True, full_cm=True):
         if not len(ds):
             ds = ["train", "valid", "test"]
         for prefix in ds:
             print(f"{prefix} dataset")
             print(f" * accuracy: {self.results[f'{prefix}_accuracy']:.3f}")
             print(f" * μ-precision: {self.results[f'{prefix}_precision']:.3f}")
             print(f" * μ-recall: {self.results[f'{prefix}_recall']:.3f}")
-            self._plot_confusion_matrix(np.array(self.results[f'{prefix}_cm']), normalize=normalize)
+            suffix = '_full' if full_cm and f'{prefix}_cm_full' in self.results else ''
+            self._plot_confusion_matrix(np.array(self.results[f'{prefix}_cm{suffix}']), normalize=normalize)
 
     def _plot_confusion_matrix(self, cm, normalize, fmt=None):
         if normalize:
@@ -272,7 +304,12 @@ def _plot_confusion_matrix(self, cm, normalize, fmt=None):
             cm = cm / s
         if fmt is None:
             fmt = "0.2f" if normalize else "d"
-        target_names = ["OTHER", "DATASET", "MODEL (paper)", "MODEL (comp.)", "METRIC", "EMPTY"]
+
+        if len(cm) == 6:
+            target_names = ["OTHER", "DATASET", "MODEL (paper)", "MODEL (comp.)", "METRIC", "EMPTY"]
+        else:
+            target_names = ["OTHER", "params", "task", "DATASET", "subdataset", "MODEL (paper)", "model (best)",
+                            "model (ens.)", "MODEL (comp.)", "METRIC", "EMPTY"]
         df_cm = pd.DataFrame(cm, index=[i for i in target_names],
                              columns=[i for i in target_names])
         plt.figure(figsize=(10, 10))