From 762167f5f3875326edf230a061f3f5612e428004 Mon Sep 17 00:00:00 2001
From: wassname <wassname@users.noreply.github.com>
Date: Tue, 26 Nov 2019 15:47:42 +0800
Subject: [PATCH] draft funtions and metrics

---
 multifit/datasets/dataset.py |  1 -
 multifit/metrics.py          | 11 +++++++++++
 multifit/training.py         | 10 +++++++---
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/multifit/datasets/dataset.py b/multifit/datasets/dataset.py
index ce7857a..715295c 100644
--- a/multifit/datasets/dataset.py
+++ b/multifit/datasets/dataset.py
@@ -224,7 +224,6 @@ class ULMFiTDataset(Dataset):
         return self._vocab
 
     def load_clas_databunch(self, bs, label_cls=None, **args):
-        print('DEBUG', bs, label_cls, args)
         vocab = self._load_vocab()
 
         cls_name = "cls.cache.databunch"
diff --git a/multifit/metrics.py b/multifit/metrics.py
index 8cdbdaf..5c9b72f 100644
--- a/multifit/metrics.py
+++ b/multifit/metrics.py
@@ -13,6 +13,17 @@ def auc_roc_score_multi(input, targ):
     return torch.tensor(scores).mean()
 
 
+def fbeta_cls_n(y_pred, y_true, class_n=1, **args):
+    """F1 score of class 1, to be used with 2 classes."""
+    y_pred = torch.nn.functional.softmax(y_pred, dim=-1)
+    return fbeta(y_pred, y_true[:, None], sigmoid=False, **args)
+
+def auc_roc_score_cls_n(y_pred, y_true, class_n=1, **args):
+    """F1 score of class 1, to be used with 2 classes."""
+    y_pred = torch.nn.functional.softmax(y_pred, dim=-1)
+    return auc_roc_score(y_pred[:, class_n], y_true==class_n, **args)
+
+
 def fbeta_binary(y_pred, y_true, **args):
     return fbeta(y_pred[:, None], y_true[:, None], **args)
 
diff --git a/multifit/training.py b/multifit/training.py
index a1302dd..bd18c6e 100644
--- a/multifit/training.py
+++ b/multifit/training.py
@@ -281,6 +281,7 @@ class ULMFiTPretraining(ULMFiTTrainingCommand):
         learn.destroy()
         self.save_paramters()
         print("Language model saved to", self.experiment_path)
+        return learn
 
     def validate(self):
         raise NotImplementedError("The validation on the language model is not implemented.")
@@ -400,19 +401,22 @@ class ULMFiTClassifier(ULMFiTTrainingCommand):
         return learn
 
 
-    def train_(self, dataset_or_path=None, label_cls=None, loss_func=None, label_cols=None, metrics=[accuracy], **train_config):
+    def train_(self, dataset_or_path=None, label_cls=None, loss_func=None, classes=None, label_cols=None, metrics=[accuracy], **train_config):
         self.replace_(**train_config, _strict=True)
 
         base_tokenizer = self.base.tokenizer
         dataset = self._set_dataset_(dataset_or_path, base_tokenizer)
-        data_clas = dataset.load_clas_databunch(bs=self.bs, label_cls=label_cls, label_cols=label_cols)
+        data_clas = dataset.load_clas_databunch(bs=self.bs, classes=classes, label_cls=label_cls, label_cols=label_cols)
         learn = self.get_learner(data_clas=data_clas, loss_func=loss_func)
         learn.metrics = metrics
+        print("RNN head", learn.model[-1].layers[-1])
+        print("label class", data_clas.test_ds.y)
+        print("dataset loss func", data_clas.loss_func)
         print(f"Training: {learn.path / learn.model_dir}")
         learn.unfreeze()
+        self.experiment_path = learn.path / learn.model_dir
         self._fit_schedule(learn)
 
-        self.experiment_path = learn.path / learn.model_dir
         base_tokenizer.save(self.experiment_path, learn=learn)
         learn.to_fp32()
         learn.save(CLS_BEST, with_opt=False)