Add ULMFiTExperiment

* add ULMFiTExperiment
* add caching to evidence extraction
* move older Experiment params to NBSVMExperiment

Author: mkardas

sota_extractor2/config.py

@@ -1,5 +1,5 @@
 import logging
-from pathlib import  Path
+from pathlib import Path
 
 logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
                     datefmt='%m/%d/%Y %H:%M:%S',

sota_extractor2/data/structure.py

@@ -46,35 +46,58 @@ def empty_fragment(paper_id):
     return fragment
 
 
-def fetch_evidence(cell_content, cell_reference, paper_id, table_name, row, col, paper_limit=10, corpus_limit=10):
+def normalize_query(query):
+    if isinstance(query, list):
+        return tuple(normalize_query(x) for x in query)
+    if isinstance(query, dict):
+        return tuple([(normalize_query(k), normalize_query(v)) for k,v in query.items()])
+    return query
+
+_evidence_cache = {}
+_cache_miss = 0
+_cache_hit = 0
+def get_cached_or_execute(query):
+    global _evidence_cache, _cache_hit, _cache_miss
+    n = normalize_query(query.to_dict())
+    if n not in _evidence_cache:
+        _evidence_cache[n] = list(query)
+        _cache_miss += 1
+    else:
+        _cache_hit += 1
+    return _evidence_cache[n]
+
+
+def fetch_evidence(cell_content, cell_reference, paper_id, table_name, row, col, paper_limit=10, corpus_limit=10,
+                   cache=False):
     if not filter_cells(cell_content):
         return [empty_fragment(paper_id)]
     cell_content = clear_cell(cell_content)
     if cell_content == "" and cell_reference == "":
         return [empty_fragment(paper_id)]
 
+    cached_query = get_cached_or_execute if cache else lambda x: x
     evidence_query = Fragment.search().highlight(
         'text', pre_tags="<b>", post_tags="</b>", fragment_size=400)
     cell_content = cell_content.replace("\xa0", " ")
     query = {
         "query": cell_content,
         "slop": 2
     }
-    paper_fragments = list(evidence_query
+    paper_fragments = list(cached_query(evidence_query
                            .filter('term', paper_id=paper_id)
-                           .query('match_phrase', text=query)[:paper_limit])
+                           .query('match_phrase', text=query)[:paper_limit]))
     if cell_reference != "":
-        reference_fragments = list(evidence_query
+        reference_fragments = list(cached_query(evidence_query
                                    .filter('term', paper_id=paper_id)
                                    .query('match_phrase', text={
                                         "query": cell_reference,
                                         "slop": 1
-                                    })[:paper_limit])
+                                    })[:paper_limit]))
     else:
         reference_fragments = []
-    other_fagements = list(evidence_query
+    other_fagements = list(cached_query(evidence_query
                            .exclude('term', paper_id=paper_id)
-                           .query('match_phrase', text=query)[:corpus_limit])
+                           .query('match_phrase', text=query)[:corpus_limit]))
 
     ext_id = f"{paper_id}/{table_name}/{row}.{col}"
     ####print(f"{ext_id} |{cell_content}|: {len(paper_fragments)} paper fragments, {len(reference_fragments)} reference fragments, {len(other_fagements)} other fragments")
@@ -137,22 +160,23 @@ def filter_cells(cell_content):
 interesting_types = ["model-paper", "model-best", "model-competing", "dataset", "dataset-sub",  "dataset-task"]
 
 
-def evidence_for_table(paper_id, table, paper_limit, corpus_limit):
+def evidence_for_table(paper_id, table, paper_limit, corpus_limit, cache=False):
     records = [
         record
             for cell in consume_cells(table)
             for evidence in fetch_evidence(cell.vals[0], cell.vals[2], paper_id=paper_id, table_name=table.name,
-                                           row=cell.row, col=cell.col, paper_limit=paper_limit, corpus_limit=corpus_limit)
+                                           row=cell.row, col=cell.col, paper_limit=paper_limit, corpus_limit=corpus_limit,
+                                           cache=cache)
             for record in create_evidence_records(evidence, cell, paper_id=paper_id, table=table)
     ]
     df = pd.DataFrame.from_records(records, columns=evidence_columns)
     return df
 
 
-def prepare_data(tables, csv_path):
+def prepare_data(tables, csv_path, cache=False):
     data = [evidence_for_table(table.paper_id, table,
                                        paper_limit=100,
-                                       corpus_limit=20) for table in progress_bar(tables)]
+                                       corpus_limit=20, cache=cache) for table in progress_bar(tables)]
     if len(data):
         df = pd.concat(data)
     else:

sota_extractor2/models/structure/experiment.py

@@ -34,9 +34,21 @@ class Labels(Enum):
 # put here to avoid recompiling, used only in _limit_context
 elastic_tag_split_re = re.compile("(<b>.*?</b>)")
 
+# e = Experiment(remove_num=False, drop_duplicates=False, vectorizer='count',
+#                this_paper=True, merge_fragments=True, merge_type='concat',
+#                evidence_source='text_highlited', split_btags=True, fixed_tokenizer=True,
+#                fixed_this_paper=True, mask=False, evidence_limit=None, context_tokens=None,
+#                analyzer='word', lowercase=True, class_weight='balanced', multinomial_type='multinomial',
+#                solver='lbfgs', C=0.1, dual=False, penalty='l2', ngram_range=[1, 3],
+#                min_df=10, max_df=0.9, max_iter=1000, results={}, has_model=False)
+
+# ULMFiT related parameters
+# remove_num, drop_duplicates, this_paper, merge_fragments, merge_type, evidence_source, split_btags
+# fixed_tokenizer?, fixed_this_paper (remove), mask, evidence_limit, context_tokens, lowercase
+# class_weight? (consider adding support),
+
 @dataclass
 class Experiment:
-    vectorizer: str = "tfidf"
     this_paper: bool = False
     merge_fragments: bool = False
     merge_type: str = "concat"  # "concat", "vote_maj", "vote_avg", "vote_max"
@@ -47,23 +59,11 @@ class Experiment:
     mask: bool = False             # if True and evidence_source = "text_highlited", replace <b>...</b> with xxmask
     evidence_limit: int = None     # maximum number of evidences per cell (grouped by (ext_id, this_paper))
     context_tokens: int = None      # max. number of words before <b> and after </b>
-    analyzer: str = "word"            # "char", "word" or "char_wb"
     lowercase: bool = True
     remove_num: bool = True
     drop_duplicates: bool = True
     mark_this_paper: bool = False
 
-    class_weight: str = None
-    multinomial_type: str = "manual"  # "manual", "ovr", "multinomial"
-    solver: str = "liblinear"  # 'lbfgs' - large, liblinear for small datasets
-    C: float = 4.0
-    dual: bool = True
-    penalty: str = "l2"
-    ngram_range: tuple = (1, 2)
-    min_df: int = 3
-    max_df: float = 0.9
-    max_iter: int = 1000
-
     results: dict = dataclasses.field(default_factory=dict)
 
     has_model: bool = False     # either there's already pretrained model or it's a saved experiment and there's a saved model as well
@@ -78,29 +78,39 @@ def _get_next_exp_name(self, dir_path):
                 return dir_path / name
         raise Exception("You have too many files in this dir, really!")
 
-    def _save_model(self, path):
+    @staticmethod
+    def _dump_pickle(obj, path):
         with open(path, 'wb') as f:
-            pickle.dump(self._model, f)
+            pickle.dump(obj, f)
 
-    def _load_model(self, path):
+    @staticmethod
+    def _load_pickle(path):
         with open(path, 'rb') as f:
-            self._model = pickle.load(f)
-            return self._model
+            return pickle.load(f)
+
+    def _save_model(self, path):
+        self._dump_pickle(self._model, path)
+
+    def _load_model(self, path):
+        self._model = self._load_pickle(path)
+        return self._model
 
     def load_model(self):
         path = self._path.parent / f"{self._path.stem}.model"
         return self._load_model(path)
 
+    def save_model(self, path):
+        if hasattr(self, "_model"):
+            self._save_model(path)
+
     def save(self, dir_path):
         dir_path = Path(dir_path)
         dir_path.mkdir(exist_ok=True, parents=True)
         filename = self._get_next_exp_name(dir_path)
         j = dataclasses.asdict(self)
         with open(filename, "wt") as f:
             json.dump(j, f)
-        if hasattr(self, "_model"):
-            fn = filename.stem
-            self._save_model(dir_path / f"{fn}.model")
+        self.save_model(dir_path / f"{filename.stem}.model")
         return filename.name
 
     def to_df(self):
@@ -119,12 +129,13 @@ def new_experiment(self, **kwargs):
     def update_results(self, **kwargs):
         self.results.update(**kwargs)
 
-    def get_trained_model(self, train_df):
-        nbsvm = NBSVM(experiment=self)
-        nbsvm.fit(train_df["text"], train_df["label"])
-        self._model = nbsvm
+    def train_model(self, train_df, valid_df):
+        raise NotImplementedError("train_model should be implemented in subclass")
+
+    def get_trained_model(self, train_df, valid_df):
+        self._model = self.train_model(train_df, valid_df)
         self.has_model = True
-        return nbsvm
+        return self._model
 
     def _limit_context(self, text):
         parts = elastic_tag_split_re.split(text)
@@ -301,3 +312,23 @@ def experiments_to_df(cls, exps):
         dfs = [e.to_df() for e in exps]
         df = pd.concat(dfs)
         return df
+
+@dataclass
+class NBSVMExperiment(Experiment):
+    vectorizer: str = "tfidf"
+    analyzer: str = "word"            # "char", "word" or "char_wb"
+    class_weight: str = None
+    multinomial_type: str = "manual"  # "manual", "ovr", "multinomial"
+    solver: str = "liblinear"  # 'lbfgs' - large, liblinear for small datasets
+    C: float = 4.0
+    dual: bool = True
+    penalty: str = "l2"
+    ngram_range: tuple = (1, 2)
+    min_df: int = 3
+    max_df: float = 0.9
+    max_iter: int = 1000
+
+    def train_model(self, train_df, valid_df=None):
+        nbsvm = NBSVM(experiment=self)
+        nbsvm.fit(train_df["text"], train_df["label"])
+        return nbsvm

sota_extractor2/models/structure/structure_predictor.py

@@ -3,7 +3,8 @@
 import pandas as pd
 import numpy as np
 import pickle
-from .experiment import Experiment, Labels, label_map
+from .experiment import Labels, label_map
+from .ulmfit_experiment import ULMFiTExperiment
 import re
 from .ulmfit import ULMFiT_SP
 from ...pipeline_logger import pipeline_logger
@@ -45,13 +46,11 @@ def __init__(self, path, file, crf_path=None, crf_model="crf.pkl",
         self.crf = load_crf(crf_path / crf_model)
 
         # todo: clean Experiment from older approaches
-        self._e = Experiment(remove_num=False, drop_duplicates=False, vectorizer='count',
-                       this_paper=True, merge_fragments=True, merge_type='concat',
-                       evidence_source='text_highlited', split_btags=True, fixed_tokenizer=True,
-                       fixed_this_paper=True, mask=False, evidence_limit=None, context_tokens=None,
-                       analyzer='word', lowercase=True, class_weight='balanced', multinomial_type='multinomial',
-                       solver='lbfgs', C=0.1, dual=False, penalty='l2', ngram_range=[1, 3],
-                       min_df=10, max_df=0.9, max_iter=1000, results={}, has_model=False)
+        self._e = ULMFiTExperiment(remove_num=False, drop_duplicates=False,
+               this_paper=True, merge_fragments=True, merge_type='concat',
+               evidence_source='text_highlited', split_btags=True, fixed_tokenizer=True,
+               fixed_this_paper=True, mask=False, evidence_limit=None, context_tokens=None,
+               lowercase=True)
 
     def preprocess_df(self, raw_df):
         return self._e.transform_df(raw_df)
@@ -140,7 +139,7 @@ def merge_all_with_preds(self, df, df_num, preds):
         df2.label = n_classes
         return df1.append(df2, ignore_index=True)
 
-
+    # todo: fix numeric cells being labelled as meta / other
     def format_predictions(self, tables_preds, test_ids):
         num2label = {v: k for k, v in label_map.items()}
         num2label[0] = "table-meta"
@@ -172,6 +171,7 @@ def label_table(self, paper, table, annotations, in_place):
         ext_id = (paper.paper_id, table.name)
         if ext_id in annotations:
             for _, entry in annotations[ext_id].iterrows():
+                # todo: add model-ensemble support
                 structure.iloc[entry.row, entry.col] = entry.predicted_tags if entry.predicted_tags != "model-paper" else "model-best"
         if not in_place:
             table = deepcopy(table)