Add experiments funcitonality

Author: Marcin Kardas

sota_extractor2/helpers/training.py

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

sota_extractor2/models/structure/__init__.py

@@ -5,7 +5,7 @@
 from ... import config
 
 def split_by_cell_content(df, seed=42, split_column="cell_content"):
-    set_seed(seed, "val_split")
+    set_seed(seed, "val_split", quiet=True)
     contents = np.random.permutation(df[split_column].unique())
     val_split = int(len(contents)*0.1)
     val_keys = contents[:val_split]

sota_extractor2/models/structure/experiment.py

@@ -0,0 +1,198 @@
+import dataclasses
+from dataclasses import dataclass
+import json
+from pathlib import Path
+import numpy as np
+import pandas as pd
+from sota_extractor2.models.structure.nbsvm import *
+from sklearn.metrics import confusion_matrix
+from matplotlib  import pyplot as plt
+import seaborn as sn
+from enum import Enum
+import pickle
+
+class Labels(Enum):
+    OTHER=0
+    DATASET=1
+    PAPER_MODEL=2
+    COMPETING_MODEL=3
+
+label_map = {
+    "dataset": Labels.DATASET.value,
+    "dataset-sub": Labels.DATASET.value,
+    "model-paper": Labels.PAPER_MODEL.value,
+    "model-best": Labels.PAPER_MODEL.value,
+    "model-competing": Labels.COMPETING_MODEL.value
+}
+
+@dataclass
+class Experiment:
+    vectorizer: str = "tfidf"
+    this_paper: bool = False
+    merge_fragments: bool = False
+    evidence_source: str = "text"  # "text" or "text_highlited"
+    split_btags: bool = False  # <b>Test</b> -> <b> Test </b>
+    fixed_tokenizer: bool = False  # <b> and </b> are not split
+
+    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
+    name: str = None
+
+    def _get_next_exp_name(self, dir_path):
+        dir_path = Path(dir_path)
+        files = [f.name for f in dir_path.glob("*.exp.json")]
+        for i in range(100000):
+            name = f"{i:05d}.exp.json"
+            if name not in files:
+                return dir_path / name
+        raise Exception("You have too many files in this dir, really!")
+
+    def _save_model(self, path):
+        with open(path, 'wb') as f:
+            pickle.dump(self._model, f)
+
+    def _load_model(self, path):
+        with open(path, 'rb') as f:
+            self._model = pickle.load(f)
+            return self._model
+
+    def load_model(self):
+        path = self._path.parent / f"{self._path.stem}.model"
+        return self._load_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")
+        return filename.name
+
+    def to_df(self):
+        d = dataclasses.asdict(self)
+        res = d.pop("results")
+        d.update(res)
+        row = pd.DataFrame({k: [v] for k, v in d.items()})
+        return row
+
+    def new_experiment(self, **kwargs):
+        # reset this fields unless their provided in load()
+        kwargs.setdefault("has_model", False)
+        kwargs.setdefault("results", {})
+        return dataclasses.replace(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
+        self.has_model = True
+        return nbsvm
+
+    def _transform_df(self, df):
+        df = df[df["cell_type"] != "table-meta"]  # otherwise we get precision 0 on test set
+        if self.evidence_source != "text":
+            df = df.copy(True)
+            df["text"] = df[self.evidence_source]
+        if self.merge_fragments:
+            df = df.groupby(by=["ext_id", "cell_content", "cell_type", "this_paper"]).text.apply(
+                lambda x: "\n".join(x.values)).reset_index()
+        df = df.drop_duplicates(["text", "cell_content", "cell_type"]).fillna("")
+        if self.this_paper:
+            df = df[df.this_paper]
+        if self.split_btags:
+            df["text"] = df["text"].replace(re.compile(r"(\</?b\>)"), r" \1 ")
+        df = df.replace(re.compile(r"(xxref|xxanchor)-[\w\d-]*"), "\\1 ")
+        df = df.replace(re.compile(r"(^|[ ])\d+\.\d+(\b|%)"), " xxnum ")
+        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))
+        df["label"] = pd.Categorical(df["label"])
+        return df
+
+    def transform_df(self, *dfs):
+        return [self._transform_df(df) for df in dfs]
+
+    def evaluate(self, model, train_df, valid_df, test_df):
+        for prefix, tdf in zip(["train", "valid", "test"], [train_df, valid_df, test_df]):
+            probs = model.predict_proba(tdf["text"])
+            preds = np.argmax(probs, axis=1)
+            true_y = tdf["label"]
+
+            m = metrics(preds, tdf.label)
+            r = {}
+            r[f"{prefix}_accuracy"] = m["accuracy"]
+            r[f"{prefix}_precision"] = m["precision"]
+            r[f"{prefix}_cm"] = confusion_matrix(true_y, preds).tolist()
+            self.update_results(**r)
+
+    def show_results(self, *ds):
+        if not len(ds):
+            ds = ["train", "valid", "test"]
+        for prefix in ds:
+            print(f"{prefix} dataset")
+            print(f" * accuracy: {self.results[f'{prefix}_accuracy']}")
+            print(f" * precision: {self.results[f'{prefix}_precision']}")
+            self._plot_confusion_matrix(np.array(self.results[f'{prefix}_cm']), normalize=True)
+
+    def _plot_confusion_matrix(self, cm, normalize):
+        if normalize:
+            cm = cm / cm.sum(axis=1)[:, None]
+        target_names = ["OTHER", "DATASET", "MODEL (paper)", "MODEL (comp.)"]
+        df_cm = pd.DataFrame(cm, index=[i for i in target_names],
+                             columns=[i for i in target_names])
+        plt.figure(figsize=(10, 10))
+        ax = sn.heatmap(df_cm,
+                        annot=True,
+                        square=True,
+                        fmt="0.2f" if normalize else "d",
+                        cmap="YlGnBu",
+                        mask=cm == 0,
+                        linecolor="black",
+                        linewidths=0.01)
+        ax.set_ylabel("True")
+        ax.set_xlabel("Predicted")
+
+    @classmethod
+    def load_all(cls, dir_path):
+        dir_path = Path(dir_path)
+        return [cls.load(f) for f in dir_path.glob("*.exp.json")]
+
+    @classmethod
+    def load(cls, path):
+        # a new field added to the class should not change
+        # the default behaviour of experiment, so that we
+        # can load older experiments by setting missing fields
+        # to their default values
+        e = cls()
+        path = Path(path)
+        with open(path, "rt") as f:
+            j = json.load(f)
+        j["name"] = path.name
+        e = e.new_experiment(**j)
+        e._path = path
+        return e
+
+    @classmethod
+    def experiments_to_df(cls, exps):
+        dfs = [e.to_df() for e in exps]
+        df = pd.concat(dfs)
+        return df

sota_extractor2/models/structure/nbsvm.py

@@ -39,55 +39,115 @@ def get_number_of_classes(y):
         return y.shape[1]
 
 class NBSVM:
-    def __init__(self, solver='liblinear', dual=True, C=4, ngram_range=(1, 2)):
-        self.solver = solver  # 'lbfgs' - large, liblinear for small datasets
-        self.dual = dual
-        self.C = C
-        self.ngram_range = ngram_range
+    def __init__(self, experiment):
+        self.experiment = experiment
 
     re_tok = re.compile(f'([{string.punctuation}“”¨«»®´·º½¾¿¡§£₤‘’])')
+    re_tok_fixed = re.compile(f'([{string.punctuation}“”¨«»®´·º½¾¿¡§£₤‘’])'.replace('<', '').replace('>', '').replace('/', ''))
 
-    def tokenize(self, s): 
+    def tokenize(self, s):
         return self.re_tok.sub(r' \1 ', s).split()
 
+    def tokenize_fixed(self, s):
+        return self.re_tok_fixed.sub(r' \1 ', s).split()
+
     def pr(self, y_i, y):
         p = self.trn_term_doc[y == y_i].sum(0)
         return (p+1) / ((y == y_i).sum()+1)
 
     def get_mdl(self, y):
         y = y.values
         r = np.log(self.pr(1, y) / self.pr(0, y))
-        m = LogisticRegression(C=self.C, dual=self.dual, solver=self.solver, max_iter=1000)
+        m = LogisticRegression(C=self.experiment.C, penalty=self.experiment.penalty,
+                               dual=self.experiment.dual, solver=self.experiment.solver,
+                               max_iter=self.experiment.max_iter)
         x_nb = self.trn_term_doc.multiply(r)
         return m.fit(x_nb, y), r
 
     def bow(self, X_train):
         self.n = X_train.shape[0]
-        self.vec = TfidfVectorizer(ngram_range=self.ngram_range, tokenizer=self.tokenize,
-                                min_df=3, max_df=0.9, strip_accents='unicode', use_idf=1,
-                                smooth_idf=1, sublinear_tf=1)
+
+        if self.experiment.vectorizer == "tfidf":
+            self.vec = TfidfVectorizer(ngram_range=self.experiment.ngram_range,
+                                       tokenizer=self.tokenize_fixed if self.experiment.fixed_tokenizer else self.tokenize,
+                                       min_df=self.experiment.min_df, max_df=self.experiment.max_df,
+                                       strip_accents='unicode', use_idf=1,
+                                       smooth_idf=1, sublinear_tf=1)
+        elif self.experiment.vectorizer == "count":
+            self.vec = CountVectorizer(ngram_range=self.experiment.ngram_range, tokenizer=self.tokenize,
+                                       min_df=self.experiment.min_df, max_df=self.experiment.max_df,
+                                       strip_accents='unicode')
+        else:
+            raise Exception(f"Unknown vectorizer type: {self.experiment.vectorizer}")
+
         return self.vec.fit_transform(X_train)
 
     def train_models(self, y_train):
         self.models = []
-        for i in range(0, self.c):
-            print('fit', i)
-            m, r = self.get_mdl(get_class_column(y_train, i))
-            self.models.append((m, r))
+        if self.experiment.multinomial_type == "manual":
+            for i in range(0, self.c):
+                #print('fit', i)
+                m, r = self.get_mdl(get_class_column(y_train, i))
+                self.models.append((m, r))
+        elif self.experiment.multinomial_type == "multinomial":
+            m = LogisticRegression(C=self.experiment.C, penalty=self.experiment.penalty,
+                                   dual=self.experiment.dual, solver=self.experiment.solver,
+                                   max_iter=self.experiment.max_iter,
+                                   multi_class="multinomial", class_weight=self.experiment.class_weight)
+            x_nb = self.trn_term_doc
+            self.models.append(m.fit(x_nb, y_train))
+        else:
+            raise Exception(f"Unsupported multinomial_type {self.experiment.multinomial_type}")
 
     def fit(self, X_train, y_train):
         self.trn_term_doc = self.bow(X_train)
         self.c = get_number_of_classes(y_train)
         self.train_models(y_train)
 
     def predict_proba(self, X_test):
-        preds = np.zeros((len(X_test), self.c))
         test_term_doc = self.vec.transform(X_test)
-        for i in range(0, self.c):
-            m, r = self.models[i]
-            preds[:, i] = m.predict_proba(test_term_doc.multiply(r))[:, 1]
+        if self.experiment.multinomial_type == "manual":
+            preds = np.zeros((len(X_test), self.c))
+            for i in range(0, self.c):
+                m, r = self.models[i]
+                preds[:, i] = m.predict_proba(test_term_doc.multiply(r))[:, 1]
+        elif self.experiment.multinomial_type == "multinomial":
+            preds = self.models[0].predict_proba(test_term_doc)
+        else:
+            raise Exception(f"Unsupported multinomial_type {self.experiment.multinomial_type}")
         return preds
-    
+
+    def sort_features_by_importance(self, label):
+        label = label.value
+        names = np.array(self.vec.get_feature_names())
+        if self.experiment.multinomial_type == "manual":
+            m, r = self.models[label]
+            f = m.coef_[0] * np.array(r[0])
+        elif self.experiment.multinomial_type == "multinomial":
+            f = self.models[0].coef_[label]
+        else:
+            raise Exception(f"Unsupported multinomial_type {self.experiment.multinomial_type}")
+        if self.experiment.vectorizer == "tfidf":
+            f *= self.vec.idf_
+        indices = f.argsort()[::-1]
+        return names[indices], f[indices]
+
+    def get_mismatched(self, df, true_label, predicted_label):
+        true_label = true_label.value
+        predicted_label = predicted_label.value
+
+        probs = self.predict_proba(df["text"])
+        preds = np.argmax(probs, axis=1)
+        true_y = df["label"]
+
+        mismatched_indices = (true_y == true_label) & (preds == predicted_label)
+        mismatched = df[mismatched_indices]
+        diff = probs[mismatched_indices, true_label] - probs[mismatched_indices, predicted_label]
+        indices = diff.argsort()
+        mismatched = mismatched.iloc[indices]
+        mismatched["pr_diff"] = diff[indices]
+        return mismatched
+
     def validate(self, X_test, y_test):
         acc = (np.argmax(self.predict_proba(X_test),  axis=1) == y_test).mean()
         return acc
@@ -98,10 +158,14 @@ def metrics(preds, true_y):
     acc = (p == y).mean()
     tp = ((y != 0) & (p == y)).sum()
     fp = ((p != 0) & (p != y)).sum()
+    fn = ((y != 0) & (p == 0)).sum()
+
     prec = tp / (fp + tp)
+    reca = tp / (fn + tp)
     return {
         "precision": prec,
         "accuracy": acc,
+        "recall": reca,
         "TP": tp,
         "FP": fp,
     }
@@ -130,6 +194,18 @@ def preds_for_cell_content_multi(test_df, probs, group_by=["cell_content"]):
                             'counts': grouped_counts})
     return results
 
+def preds_for_cell_content_best(test_df, probs, group_by=["cell_content"]):
+    test_df = test_df.copy()
+    probs_df = pd.DataFrame(probs, index=test_df.index)
+    test_df = pd.concat([test_df, probs_df], axis=1)
+    grouped_preds = np.argmax(test_df.groupby(
+        group_by)[probs_df.columns].sum().values, axis=1)
+    grouped_counts = test_df.groupby(group_by)["label"].count()
+    results = pd.DataFrame({'true': test_df.groupby(group_by)["label"].agg(lambda x: x.value_counts().index[0]),
+                            'pred': grouped_preds,
+                            'counts': grouped_counts})
+    return results
+
 def test_model(model, tdf):
     probs = model(tdf["text"])
     preds = np.argmax(probs, axis=1)