Allow disabling CRF in structure prediction

Author: mkardas

sota_extractor2/models/structure/structure_predictor.py

@@ -66,16 +66,20 @@ def df2tl(self, df):
         df = df[text_cols]
         return TextList.from_df(df, cols=text_cols)
 
-    def get_features(self, evidences):
+    def get_features(self, evidences, use_crf=True):
+        if use_crf:
+            learner = self.learner
+        else:
+            learner = self._full_learner
         if len(evidences):
             tl = self.df2tl(evidences)
-            self.learner.data.add_test(tl)
+            learner.data.add_test(tl)
 
-            preds, _ = self.learner.get_preds(DatasetType.Test, ordered=True)
+            preds, _ = learner.get_preds(DatasetType.Test, ordered=True)
             return preds.cpu().numpy()
-        return np.zeros((0, n_ulmfit_features))
+        return np.zeros((0, n_ulmfit_features if use_crf else n_classes))
 
-    def to_tables(self, df, transpose=False):
+    def to_tables(self, df, transpose=False, n_ulmfit_features=n_ulmfit_features):
         X_tables = []
         Y_tables = []
         ids = []
@@ -127,12 +131,12 @@ def merge_with_preds(self, df, preds):
         return list(zip(ext_id[0] + "/" + ext_id[1], ext_id[2].astype(int), ext_id[3].astype(int),
                         preds, df.text, df.cell_content, df.cell_layout, df.cell_styles, df.cell_reference, df.label))
 
-    def merge_all_with_preds(self, df, df_num, preds):
+    def merge_all_with_preds(self, df, df_num, preds, use_crf=True):
         columns = ["table_id", "row", "col", "features", "text", "cell_content", "cell_layout",
                    "cell_styles", "cell_reference", "label"]
 
         alpha = self.merge_with_preds(df, preds)
-        nums = self.merge_with_preds(df_num, np.zeros((len(df_num), n_ulmfit_features)))
+        nums = self.merge_with_preds(df_num, np.zeros((len(df_num), n_ulmfit_features if use_crf else n_classes)))
 
         df1 = pd.DataFrame(alpha, columns=columns)
         df2 = pd.DataFrame(nums, columns=columns)
@@ -156,13 +160,16 @@ def format_predictions(self, tables_preds, test_ids):
                  labels[r, c]])
         return pd.DataFrame(flat, columns=["paper", "table", "row", "col", "predicted_tags"])
 
-    def predict_tags(self, raw_evidences):
+    def predict_tags(self, raw_evidences, use_crf=True):
         evidences, evidences_num = self.keep_alphacells(self.preprocess_df(raw_evidences))
         pipeline_logger(f"{TableStructurePredictor.step}::evidences_split", evidences=evidences, evidences_num=evidences_num)
-        features = self.get_features(evidences)
-        df = self.merge_all_with_preds(evidences, evidences_num, features)
-        tables, contents, ids = self.to_tables(df)
-        preds = self.crf.predict(tables)
+        features = self.get_features(evidences, use_crf)
+        df = self.merge_all_with_preds(evidences, evidences_num, features, use_crf)
+        tables, contents, ids = self.to_tables(df, n_ulmfit_features=n_ulmfit_features if use_crf else n_classes)
+        if use_crf:
+            preds = self.crf.predict(tables)
+        else:
+            preds = [table[..., :n_classes].argmax(axis=-1) for table in tables]
         return self.format_predictions(preds, ids)
 
     # todo: consider adding sota/ablation information
@@ -179,10 +186,10 @@ def label_table(self, paper, table, annotations, in_place):
         return table
 
     # todo: take EvidenceExtractor in constructor
-    def label_tables(self, paper, tables, raw_evidences, in_place=False):
+    def label_tables(self, paper, tables, raw_evidences, in_place=False, use_crf=True):
         pipeline_logger(f"{TableStructurePredictor.step}::label_tables", paper=paper, tables=tables, raw_evidences=raw_evidences)
         if len(raw_evidences):
-            tags = self.predict_tags(raw_evidences)
+            tags = self.predict_tags(raw_evidences, use_crf)
             annotations = dict(list(tags.groupby(by=["paper", "table"])))
         else:
             annotations = {}  # just deep-copy all tables