Add ensemble commnad

ulmfit/__main__.py

@@ -1,16 +1,14 @@
 import gc
 import os
-import re
-import pprint
 import tarfile
-import shutil
-from collections import OrderedDict
+from collections import defaultdict
+import re
 from functools import wraps
 import numpy as np
 import pandas as pd
 import fire
 
-from .pretrain_lm import LMHyperParams
+from .pretrain_lm import LMHyperParams, folder_name_to_model_name, DataSetParams
 from .train_clas import CLSHyperParams
 from pathlib import Path
 from string import Template
@@ -33,10 +31,6 @@ def get_dataset_path(p, dataset_template):
     for ds_path in ds.parent.glob(pattern):
         yield lang, ds_path
 
-name_re = re.compile("(bwd)?(lstm|qrnn)_(.*)_(lmseed-)?.*\.m")
-def folder_name_to_model_name(folder_name):
-    return name_re.match(folder_name).group(3)
-
 class ULMFiT:
     @wraps(LMHyperParams)
     def lm(self, dataset_path,  **changes):
@@ -156,6 +150,55 @@ class ULMFiT:
                 results.append((base_model, lang, dataset_path))
         return results
 
+    # file_glob = "${ds_name}/${lang}.train.csv"
+    def ensemble(self, glob="data/mldoc*/*-1/models/sp15k/qrnn_*.m",
+                 file_template="${model_dir}/preds-on-test.npy",
+                 gold_labels_template="${dataset_path}/${lang}.test.csv",
+                 out_template="${key}.ensemble.csv",
+                 key_template='${lang}', verbose=False, exclude_re=None):
+
+        def load_labels(file, verbose=False):
+            if file.suffix == ".npy":
+                labels = np.load(str(file))
+
+            elif file.suffix == ".csv":
+                df = pd.read_csv(file, header=None)
+                labels =  np.array([df[c] for c in df.columns if np.issubdtype(df[c].dtype, np.number)]).T.squeeze()
+            else:
+                raise AttributeError("Unknown result file type", file.extension)
+            if verbose: print(file, labels.shape)
+            return labels
+
+        files_for_ensemble = defaultdict(list)
+        gold_labels = {}
+        for folder in Path.cwd().glob(glob):
+            if exclude_re is not None and re.match(exclude_re, str(folder)):
+                print("Skipping", folder)
+                continue
+            if folder.suffix == ".m":
+                params = CLSHyperParams.from_json(folder)
+            else:
+                params = DataSetParams(folder)
+
+            key = params.resolve_template(key_template)
+            file_glob = params.resolve_template(file_template)
+            files_for_ensemble[key].append(Path(file_glob))
+            gold_label_glob = params.resolve_template(gold_labels_template)
+            gold_file = Path(gold_label_glob)
+            gold_labels[key] = gold_file
+
+        for key, files in files_for_ensemble.items():
+            ensemble = np.array([load_labels(file, verbose) for file in files]).mean(axis=0)
+            if len(ensemble.shape) != 1:
+                ensemble = np.argmax(ensemble, axis=1)
+            test = pd.read_csv(gold_labels[key], header=None)
+            print({"Key": key, "Test Accuracy": (test[0] == ensemble).mean(), "on": gold_labels[key], 'files_count':len(files)})
+            test[0] = ensemble
+            if out_template:
+                out_file = Template(out_template).substitute(key=key)
+                test.to_csv(out_file, header=None)
+                print({"File saved to": out_file})
+
     def eval(self, glob="data/mldoc/*-1/models/sp30k/lstm_nl4.m", dataset_template='${ds_name}', name=None,
              num_lm_epochs=0, train=True, to_csv=None, return_df=False, label_smoothing_eps=0.0,
              lmseed=None, ftseed=None, clsweightseed=None, clstrainseed=None, save_name="cls_best",

ulmfit/pretrain_lm.py

@@ -4,7 +4,8 @@ expected to have been tokenized with Moses and processed with `postprocess_wikit
 That is, the data is expected to be white-space separated and numbers are expected
 to be split.
 """
-from dataclasses import InitVar
+from dataclasses import InitVar, asdict
+from string import Template
 
 import fastai
 import fire
@@ -56,10 +57,35 @@ def read_wiki_articles(filename):
     print(f"Wiki text was split to {len(articles)} articles")
     return pd.DataFrame({'texts': np.array(articles, dtype=np.object)})
 
-@dataclass
-class LMHyperParams:
-    dataset_path: str # data_dir
+name_re = re.compile("(bwd)?(lstm|qrnn)_(.*)_(lmseed-)?.*\.m")
+def folder_name_to_model_name(folder_name):
+    if hasattr(folder_name, 'name'):
+        folder_name = folder_name.name
+    match = name_re.match(folder_name)
+    if match:
+        return match.group(3)
+    return None
 
+@dataclass
+class DataSetParams:
+    dataset_path: str  # data_dir
+    lang: str = None
+
+    def __post_init__(self):
+        if self.lang is None:
+            self.lang = infer_lang_from_dataset(self.dataset_path.name)
+
+    def resolve_template(self, template, **additional_options):
+        try:
+            params = asdict(self)
+            params.update(additional_options)
+            params["dataset_name"] = self.dataset_path.name
+            return Template(template).substitute(**params)
+        except KeyError as e:
+            raise KeyError(f"{e} , options:{repr(list(params.keys()))}")
+
+@dataclass
+class LMHyperParams(DataSetParams):
     base_lm_path: str = None
     backwards: str = False
     bidir: bool =False
@@ -85,7 +111,6 @@ class LMHyperParams:
     rnn_alpha: float = 2  # activation regularization (AR)
     rnn_beta: float = 1  # temporal activation regularization (TAR)
 
-    lang: str = 'en'
     name: str = None
     cuda_id: InitVar[int] = 0
 
@@ -104,9 +129,6 @@ class LMHyperParams:
         self.cache_dir = self.dataset_path / 'models' / self.tokenizer_prefix
         self.model_dir = self.cache_dir / self.model_name
 
-        print('Max vocab:', self.max_vocab)
-        print('Cache dir:', self.cache_dir)
-        print('Model dir:', self.model_dir)
         if self.nh is None: self.nh = 1550 if self.qrnn else 1150
         if self.name is None: self.name = self.lang
 
@@ -191,6 +213,10 @@ class LMHyperParams:
         print("Saving info", self.model_dir / 'info.json')
 
     def train_lm(self, num_epochs=20, data_lm=None, bs=70, true_wd=False, drop_mult=0.0, lr=5e-3, label_smoothing_eps=0.0):
+        print("Training lm")
+        print('Max vocab:', self.max_vocab)
+        print('Cache dir:', self.cache_dir)
+        print('Model dir:', self.model_dir)
         if self.pretrained_fnames or self.pretrained_model:
             self.set_seed(self.ftseed, "fine-tune")
         else:
@@ -342,10 +368,11 @@ class LMHyperParams:
 
         d.update(kwargs)
         return cls(**d)
+
     @classmethod
     def from_json(cls, model_path:Path, **kwargs):
         model_path = Path(model_path).resolve()
-        name = re.search(r"[a-z]+_(.+).m", model_path.name).group(1)
+        name = folder_name_to_model_name(model_path)
         with open(model_path / 'info.json', 'r') as f:
             d = json.load(f)
         d.update(kwargs)
@@ -355,6 +382,9 @@ class LMHyperParams:
         d['lang'] = infer_lang_from_dataset(dataset_path.name)
         return cls(**d)
 
+    def resolve_template(self, template, **additional_options):
+        return super().resolve_template(template, model_dir=self.model_dir, **additional_options)
+
 def infer_lang_from_dataset(name:str):
     return name.split("-")[0]
 

ulmfit/train_clas.py

@@ -2,6 +2,7 @@
 Train a classifier on top of a language model trained with `pretrain_lm.py`.
 Optionally fine-tune LM before.
 """
+import re
 
 from fastai.callbacks import CSVLogger, SaveModelCallback
 from fastai.text import *
@@ -125,6 +126,10 @@ class CLSHyperParams(LMHyperParams):
     def train_cls(self, num_lm_epochs, unfreeze=True, num_cls_frozen_epochs=1, bs=40, drop_mul_lm=0.3, drop_mul_cls=0.5,
                   use_test_for_validation=False, num_cls_epochs=2, limit=None, noise=0.0, cls_max_len=20*70, lr_sched='layered',
                   label_smoothing_eps=0.0, random_init=False, dump_preds=None, early_stopping=True, weighted_cross_entropy=True):
+        print("Training CLS")
+        print('Max vocab:', self.max_vocab)
+        print('Cache dir:', self.cache_dir)
+        print('Model dir:', self.model_dir)
         assert use_test_for_validation == False, "use_test_for_validation=True is not supported"
         self.model_dir.mkdir(exist_ok=True, parents=True)
 
@@ -330,7 +335,4 @@ class CLSHyperParams(LMHyperParams):
         return self.databunch(name, bunch_class=TextClasDataBunch, *args, **kwargs)
 
 if __name__ == '__main__':
-    fire.Fire(CLSHyperParams)
-
-##
-
+    fire.Fire(CLSHyperParams)