Upload HGTrainer.py

HGTrainer.py

@@ -0,0 +1,432 @@
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple, Literal
+
+import torch
+import numpy
+
+from transformers import Trainer, PreTrainedModel, RobertaForSequenceClassification, BatchEncoding, RobertaConfig, \
+    EvalPrediction
+from transformers.modeling_outputs import SequenceClassifierOutput, BaseModelOutput
+from loguru import logger
+
+
+def val_nov_loss(is_val: torch.Tensor, should_val: torch.Tensor, is_nov: torch.Tensor, should_nov: torch.Tensor,
+                 weights: Optional[torch.Tensor] = None, reduce: bool = True) -> torch.Tensor:
+    if weights is None:
+        weights = torch.ones_like(should_val)
+        logger.debug("No weights-vector - assume, all {} samples should count equally", weights.size())
+
+    loss_validity = torch.pow(is_val - torch.where(torch.isnan(should_val), is_val, should_val), 2)
+    loss_novelty = torch.pow(is_nov - torch.where(torch.isnan(should_nov), is_nov, should_nov), 2)
+
+    logger.trace("loss_validity: {} / loss_novelty: {}", loss_validity, loss_novelty)
+
+    loss = (.5 * (loss_validity * loss_novelty) + .5 * loss_validity + .5 * loss_novelty) * weights
+
+    return torch.mean(loss) if reduce else loss
+
+
+def val_nov_metric(eval_data: EvalPrediction) -> Dict[str, float]:
+    if isinstance(eval_data.predictions, Tuple) and isinstance(eval_data.label_ids, Tuple) \
+            or min(len(eval_data.predictions), len(eval_data.label_ids)) >= 2:
+        logger.trace("Format is as processable ({}: {})", type(eval_data.predictions), len(eval_data.predictions))
+        if len(eval_data.predictions) != 2:
+            logger.debug("We expect 2 tuples, but get {}: {}", len(eval_data.predictions), eval_data.predictions)
+
+        is_validity = eval_data.predictions[-2]
+        should_validity = eval_data.label_ids[-2]
+        is_novelty = eval_data.predictions[-1]
+        should_novelty = eval_data.label_ids[-1]
+
+        return _val_nov_metric(is_validity=is_validity, should_validity=should_validity,
+                               is_novelty=is_novelty, should_novelty=should_novelty)
+    else:
+        logger.warning("This metric can't return all metrics properly, "
+                       "because validity and novelty are not distinguishable")
+
+        return {
+            "size": numpy.size(eval_data.label_ids),
+            "mse_validity": numpy.mean((eval_data.predictions-eval_data.label_ids) ** 2),
+            "mse_novelty": numpy.mean((eval_data.predictions-eval_data.label_ids) ** 2),
+            "error_validity": numpy.mean(numpy.abs(eval_data.predictions-eval_data.label_ids)),
+            "error_novelty": numpy.mean(numpy.abs(eval_data.predictions-eval_data.label_ids)),
+            "approximately_hits_validity": -1,
+            "approximately_hits_novelty": -1,
+            "exact_hits_validity": -1,
+            "exact_hits_novelty": -1,
+            "approximately_hits": numpy.count_nonzero(
+                numpy.where(numpy.abs(eval_data.predictions-eval_data.label_ids) < .2, 1, 0)
+            ) / numpy.size(eval_data.predictions),
+            "exact_hits": numpy.count_nonzero(
+                numpy.where(numpy.abs(eval_data.predictions-eval_data.label_ids) < .05, 1, 0)
+            ) / numpy.size(eval_data.predictions),
+            "accuracy_validity": -1,
+            "accuracy_novelty": -1,
+            "accuracy": -1,
+            "f1_validity": -1,
+            "f1_novelty": -1,
+            "f1_macro": -1,
+            "never_predicted_classes": 4
+        }
+
+
+def _val_nov_metric(is_validity: numpy.ndarray, should_validity: numpy.ndarray,
+                    is_novelty: numpy.ndarray, should_novelty: numpy.ndarray) -> Dict[str, float]:
+    ret = {
+        "size": numpy.size(is_validity),
+        "mse_validity": numpy.mean((is_validity - should_validity) ** 2),
+        "mse_novelty": numpy.mean((is_novelty - should_novelty) ** 2),
+        "error_validity": numpy.mean(numpy.abs(is_validity - should_validity)),
+        "error_novelty": numpy.mean(numpy.abs(is_novelty - should_novelty)),
+        "approximately_hits_validity": numpy.sum(
+            numpy.where(numpy.abs(is_validity - should_validity) < .2, 1, 0)) / numpy.size(is_validity),
+        "approximately_hits_novelty": numpy.sum(
+            numpy.where(numpy.abs(is_novelty - should_novelty) < .2, 1, 0)) / numpy.size(is_novelty),
+        "exact_hits_validity": numpy.sum(
+            numpy.where(numpy.abs(is_validity - should_validity) < .05, 1, 0)) / numpy.size(is_validity),
+        "exact_hits_novelty": numpy.sum(
+            numpy.where(numpy.abs(is_novelty - should_novelty) < .05, 1, 0)) / numpy.size(is_novelty),
+        "approximately_hits": numpy.sum(
+            numpy.where(numpy.abs(is_validity - should_validity) + numpy.abs(is_novelty - should_novelty) < .25, 1, 0)
+        ) / numpy.size(is_validity),
+        "exact_hits": numpy.sum(
+            numpy.where(numpy.abs(is_validity - should_validity) + numpy.abs(is_novelty - should_novelty) < .05, 1, 0)
+        ) / numpy.size(is_validity),
+        "accuracy_validity": numpy.sum(numpy.where(
+            numpy.any(numpy.stack([
+                numpy.all(numpy.stack([is_validity >= .5, should_validity >= .5]), axis=0),
+                numpy.all(numpy.stack([is_validity < .5, should_validity < .5]), axis=0)
+            ]), axis=0),
+            1, 0
+        )) / numpy.size(is_validity),
+        "accuracy_novelty": numpy.sum(numpy.where(
+            numpy.any(numpy.stack([
+                numpy.all(numpy.stack([is_novelty >= .5, should_novelty >= .5]), axis=0),
+                numpy.all(numpy.stack([is_novelty < .5, should_novelty < .5]), axis=0)
+            ]), axis=0),
+            1, 0
+        )) / numpy.size(is_validity),
+        "accuracy": numpy.sum(numpy.where(
+            numpy.any(numpy.stack([
+                numpy.all(numpy.stack([is_validity >= .5, should_validity >= .5, is_novelty >= .5, should_novelty >= .5]),
+                          axis=0),
+                numpy.all(numpy.stack([is_validity >= .5, should_validity >= .5, is_novelty < .5, should_novelty < .5]),
+                          axis=0),
+                numpy.all(numpy.stack([is_validity < .5, should_validity < .5, is_novelty >= .5, should_novelty >= .5]),
+                          axis=0),
+                numpy.all(numpy.stack([is_validity < .5, should_validity < .5, is_novelty < .5, should_novelty < .5]),
+                          axis=0)
+            ]), axis=0),
+            1, 0
+        )) / numpy.size(is_validity),
+        "never_predicted_classes": sum(
+            [int(numpy.all(numpy.abs(is_validity-validity) < .5) and numpy.all(numpy.abs(is_novelty-novelty) < .5))
+             for validity, novelty in [(1, 1), (1, 0), (0, 1), (0, 0)]]
+        )
+    }
+
+    ret_base_help = {
+        "true_positive_validity": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity >= .5, should_validity >= .5]), axis=0),
+            1, 0)),
+        "true_negative_validity": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity < .5, should_validity < .5]), axis=0),
+            1, 0)),
+        "true_positive_novelty": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_novelty >= .5, should_novelty >= .5]), axis=0),
+            1, 0)),
+        "true_negative_novelty": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_novelty < .5, should_novelty < .5]), axis=0),
+            1, 0)),
+        "true_positive_valid_novel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity >= .5, is_novelty >= .5,
+                                   should_validity >= .5, should_novelty >= .5]), axis=0),
+            1, 0)),
+        "true_positive_nonvalid_novel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity < .5, is_novelty >= .5,
+                                   should_validity < .5, should_novelty >= .5]), axis=0),
+            1, 0)),
+        "true_positive_valid_nonnovel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity >= .5, is_novelty < .5,
+                                   should_validity >= .5, should_novelty < .5]), axis=0),
+            1, 0)),
+        "true_positive_nonvalid_nonnovel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity < .5, is_novelty < .5,
+                                   should_validity < .5, should_novelty < .5]), axis=0),
+            1, 0)),
+        "classified_positive_validity": numpy.sum(numpy.where(is_validity >= .5, 1, 0)),
+        "classified_negative_validity": numpy.sum(numpy.where(is_validity < .5, 1, 0)),
+        "classified_positive_novelty": numpy.sum(numpy.where(is_novelty >= .5, 1, 0)),
+        "classified_negative_novelty": numpy.sum(numpy.where(is_novelty < .5, 1, 0)),
+        "classified_positive_valid_novel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity >= .5, is_novelty >= .5]), axis=0),
+            1, 0)),
+        "classified_positive_nonvalid_novel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity < .5, is_novelty >= .5]), axis=0),
+            1, 0)),
+        "classified_positive_valid_nonnovel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity >= .5, is_novelty < .5]), axis=0),
+            1, 0)),
+        "classified_positive_nonvalid_nonnovel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([is_validity < .5, is_novelty < .5]), axis=0),
+            1, 0)),
+        "indeed_positive_validity": numpy.sum(numpy.where(should_validity >= .5, 1, 0)),
+        "indeed_negative_validity": numpy.sum(numpy.where(should_validity < .5, 1, 0)),
+        "indeed_positive_novelty": numpy.sum(numpy.where(should_novelty >= .5, 1, 0)),
+        "indeed_negative_novelty": numpy.sum(numpy.where(should_novelty < .5, 1, 0)),
+        "indeed_positive_valid_novel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([should_validity >= .5, should_novelty >= .5]), axis=0),
+            1, 0)),
+        "indeed_positive_nonvalid_novel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([should_validity < .5, should_novelty >= .5]), axis=0),
+            1, 0)),
+        "indeed_positive_valid_nonnovel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([should_validity >= .5, should_novelty < .5]), axis=0),
+            1, 0)),
+        "indeed_positive_nonvalid_nonnovel": numpy.sum(numpy.where(
+            numpy.all(numpy.stack([should_validity < .5, should_novelty < .5]), axis=0),
+            1, 0)),
+    }
+
+    ret_help = {
+        "precision_validity": ret_base_help["true_positive_validity"] /
+                              max(1, ret_base_help["classified_positive_validity"]),
+        "precision_novelty": ret_base_help["true_positive_novelty"] /
+                             max(1, ret_base_help["classified_positive_novelty"]),
+        "recall_validity": ret_base_help["true_positive_validity"] /
+                           max(1, ret_base_help["indeed_positive_validity"]),
+        "recall_novelty": ret_base_help["true_positive_novelty"] /
+                          max(1, ret_base_help["indeed_positive_novelty"]),
+        "precision_val_neg": ret_base_help["true_negative_validity"] /
+                              max(1, ret_base_help["classified_negative_validity"]),
+        "precision_nov_neg": ret_base_help["true_negative_novelty"] /
+                             max(1, ret_base_help["classified_negative_novelty"]),
+        "recall_val_neg": ret_base_help["true_negative_validity"] /
+                           max(1, ret_base_help["indeed_negative_validity"]),
+        "recall_nov_neg": ret_base_help["true_negative_novelty"] /
+                          max(1, ret_base_help["indeed_negative_novelty"]),
+        "precision_valid_novel": ret_base_help["true_positive_valid_novel"] /
+                                 max(1, ret_base_help["classified_positive_valid_novel"]),
+        "precision_valid_nonnovel": ret_base_help["true_positive_valid_nonnovel"] /
+                                    max(1, ret_base_help["classified_positive_valid_nonnovel"]),
+        "precision_nonvalid_novel": ret_base_help["true_positive_nonvalid_novel"] /
+                                    max(1, ret_base_help["classified_positive_nonvalid_novel"]),
+        "precision_nonvalid_nonnovel": ret_base_help["true_positive_nonvalid_nonnovel"] /
+                                       max(1, ret_base_help["classified_positive_nonvalid_nonnovel"]),
+        "recall_valid_novel": ret_base_help["true_positive_valid_novel"] /
+                              max(1, ret_base_help["indeed_positive_valid_novel"]),
+        "recall_valid_nonnovel": ret_base_help["true_positive_valid_nonnovel"] /
+                                 max(1, ret_base_help["indeed_positive_valid_nonnovel"]),
+        "recall_nonvalid_novel": ret_base_help["true_positive_nonvalid_novel"] /
+                                 max(1, ret_base_help["indeed_positive_nonvalid_novel"]),
+        "recall_nonvalid_nonnovel": ret_base_help["true_positive_nonvalid_nonnovel"] /
+                                    max(1, ret_base_help["indeed_positive_nonvalid_nonnovel"])
+    }
+
+    ret.update({
+        "f1_validity": 2 * ret_help["precision_validity"] * ret_help["recall_validity"] /
+                       max(1e-4, ret_help["precision_validity"] + ret_help["recall_validity"]),
+        "f1_novelty": 2 * ret_help["precision_novelty"] * ret_help["recall_novelty"] /
+                      max(1e-4, ret_help["precision_novelty"] + ret_help["recall_novelty"]),
+        "f1_val_neg": 2 * ret_help["precision_val_neg"] * ret_help["recall_val_neg"] /
+                      max(1e-4, ret_help["precision_val_neg"] + ret_help["recall_val_neg"]),
+        "f1_nov_neg": 2 * ret_help["precision_nov_neg"] * ret_help["recall_nov_neg"] /
+                      max(1e-4, ret_help["precision_nov_neg"] + ret_help["recall_nov_neg"]),
+        "f1_valid_novel": 2 * ret_help["precision_valid_novel"] * ret_help["recall_valid_novel"] /
+                          max(1e-4, ret_help["precision_valid_novel"] + ret_help["recall_valid_novel"]),
+        "f1_valid_nonnovel": 2 * ret_help["precision_valid_nonnovel"] * ret_help["recall_valid_nonnovel"] /
+                             max(1e-4, ret_help["precision_valid_nonnovel"] + ret_help["recall_valid_nonnovel"]),
+        "f1_nonvalid_novel": 2 * ret_help["precision_nonvalid_novel"] * ret_help["recall_nonvalid_novel"] /
+                             max(1e-4, ret_help["precision_nonvalid_novel"] + ret_help["recall_nonvalid_novel"]),
+        "f1_nonvalid_nonnovel": 2 * ret_help["precision_nonvalid_nonnovel"] * ret_help["recall_nonvalid_nonnovel"] /
+                                max(1e-4, ret_help["precision_nonvalid_nonnovel"] + ret_help["recall_nonvalid_nonnovel"])
+    })
+
+    ret.update({
+        "f1_val_macro": (ret["f1_validity"] + ret["f1_val_neg"])/2,
+        "f1_nov_macro": (ret["f1_novelty"] + ret["f1_nov_neg"])/2,
+        "f1_macro": (ret["f1_valid_novel"]+ret["f1_valid_nonnovel"]+ret["f1_nonvalid_novel"]+ret["f1_nonvalid_nonnovel"])/4
+    })
+
+    logger.info("Clean the metric-dict before returning: {}",
+                " / ".join(map(lambda key: "{}: {}".format(key, ret.pop(key)),
+                               ["approximately_hits_validity", "approximately_hits_novelty", "exact_hits_validity",
+                                "exact_hits_novelty", "size"])))
+
+    return ret
+
+
+# noinspection PyMethodMayBeStatic
+class ValNovTrainer(Trainer):
+    def compute_loss(self, model: PreTrainedModel, inputs: Dict[str, torch.Tensor], return_outputs=False):
+        try:
+            validity = inputs.pop("validity")
+            novelty = inputs.pop("novelty")
+            weights = inputs.pop("weight")
+            logger.trace("The batch contain following validity-scores ({}), novelty-scores ({}) and weights ({})",
+                         validity, novelty, weights)
+
+            outputs = model(**inputs)
+
+            if isinstance(outputs, ValNovOutput) and outputs.loss is not None:
+                logger.debug("The loss was already computed: {}", outputs.loss)
+                return (outputs.loss, outputs) if return_outputs else outputs.loss
+
+            if isinstance(outputs, ValNovOutput):
+                is_val = outputs.validity
+                is_nov = outputs.novelty
+            else:
+                logger.warning("The output of you model {} is a {}, bit should be a ValNovOutput",
+                               model.name_or_path, type(outputs))
+                is_val = outputs[0] if isinstance(outputs, Tuple) and len(outputs) >= 2 else outputs
+                is_nov = outputs[1] if isinstance(outputs, Tuple) and len(outputs) >= 2 else outputs
+
+            loss = val_nov_loss(is_val=is_val, is_nov=is_nov,
+                                should_val=validity, should_nov=novelty,
+                                weights=weights)
+
+            return (loss, outputs) if return_outputs else loss
+        except KeyError:
+            logger.opt(exception=True).error("Something in your configuration / plugged model is false")
+
+        return (torch.zeros((0,), dtype=torch.float), model(**inputs)) if return_outputs \
+            else torch.zeros((0,), dtype=torch.float)
+
+
+@dataclass
+class ValNovOutput(SequenceClassifierOutput):
+    validity: torch.FloatTensor = None
+    novelty: torch.FloatTensor = None
+
+
+class ValNovRegressor(torch.nn.Module):
+    def __init__(self, transformer: PreTrainedModel,
+                 loss: Literal["ignore", "compute", "compute and reduce"] = "ignore"):
+        super(ValNovRegressor, self).__init__()
+
+        self.transformer = transformer
+        try:
+            self.regression_layer_validity = torch.nn.Linear(in_features=transformer.config.hidden_size, out_features=1)
+            self.regression_layer_novelty = torch.nn.Linear(in_features=transformer.config.hidden_size, out_features=1)
+        except AttributeError:
+            logger.opt(exception=True).warning("No hidden-size... please use a XXXForMaskedLM-Model!")
+            self.regression_layer_validity = torch.nn.LazyLinear(out_features=1)
+            self.regression_layer_novelty = torch.nn.LazyLinear(out_features=1)
+
+        self.sigmoid = torch.nn.Sigmoid()
+        if loss == "ignore":
+            logger.info("torch-Module without an additional loss computation during the forward-pass - "
+                        "has to be done explicitly in the training loop!")
+        self.loss = loss
+
+        logger.success("Successfully created {}", self)
+
+    def forward(self, x: BatchEncoding) -> ValNovOutput:
+        transformer_cls: BaseModelOutput = self.transformer(input_ids=x["input_ids"],
+                                                            attention_mask=x["attention_mask"],
+                                                            token_type_ids=x["token_type_ids"],
+                                                            return_dict=True)
+
+        cls_logits = transformer_cls.last_hidden_state[0]
+
+        validity_logits = self.regression_layer_validity(cls_logits)
+        novelty_logits = self.regression_layer_novelty(cls_logits)
+
+        return ValNovOutput(
+            logits=torch.stack([validity_logits, novelty_logits]),
+            loss=val_nov_loss(is_val=self.sigmoid(validity_logits),
+                              is_nov=self.sigmoid(novelty_logits),
+                              should_val=x["validity"],
+                              should_nov=x["novelty"],
+                              weights=x.get("weight", None),
+                              reduce=self.loss == "compute and reduce"
+                              ) if self.loss != "ignore" and "validity" in x and "novelty" in x else None,
+            hidden_states=transformer_cls.hidden_states,
+            attentions=transformer_cls.attentions,
+            validity=self.sigmoid(validity_logits),
+            novelty=self.sigmoid(novelty_logits)
+        )
+
+    def __str__(self) -> str:
+        return "() --> ({} --> validity/ {} --> novelty)".format(self.transformer.name_or_path,
+                                                                 self.regression_layer_validity,
+                                                                 self.regression_layer_novelty)
+
+
+class RobertaForValNovRegression(RobertaForSequenceClassification):
+    def __init__(self, *model_args, **model_kwargs):
+        config = RobertaForValNovRegression.get_config()
+
+        configs = [arg for arg in model_args if isinstance(arg, RobertaConfig)]
+        if len(configs) >= 1:
+            logger.warning("Found already {} config {}... extend it", len(configs), configs[0])
+            model_args = [arg for arg in model_args if not isinstance(arg, RobertaConfig)]
+            config = configs[0]
+            config.num_labels = 2
+            config.id2label = {
+                0: "validity",
+                1: "novelty"
+            }
+            config.return_dict = True
+
+        super().__init__(config=config, *model_args, **model_kwargs)
+
+        self.loss = "compute"
+        self.sigmoid = torch.nn.Sigmoid()
+
+    @classmethod
+    def get_config(cls) -> RobertaConfig:
+        config = RobertaConfig()
+        config.finetuning_task = "Validity-Novelty-Prediction"
+        config.num_labels = 2
+        config.id2label = {
+            0: "validity",
+            1: "novelty"
+        }
+        config.return_dict = True
+
+        return config
+
+    def forward(self, **kwargs):
+        logger.trace("Found {} forward-params", len(kwargs))
+        if "labels" in kwargs:
+            labels = kwargs.pop("labels")
+            logger.warning("Found a disturbing param in forward-function: labels ({})", labels)
+        if "return_dict" in kwargs:
+            return_dict = kwargs.pop("return_dict")
+            logger.warning("Found a disturbing param in forward-function: return_dict ({})", return_dict)
+
+        should_validity = None
+        if "validity" in kwargs:
+            should_validity = kwargs.pop("validity")
+            logger.trace("Found a target validity-vector: {}", should_validity)
+
+        should_novelty = None
+        if "novelty" in kwargs:
+            should_novelty = kwargs.pop("novelty")
+            logger.trace("Found a target novelty-vector: {}", should_novelty)
+
+        weights = None
+        if "weight" in kwargs:
+            weights = kwargs.pop("weight")
+            logger.trace("Found a sample-weights-vector: {}", weights)
+
+        out: SequenceClassifierOutput = super().forward(**kwargs)
+        is_validity = self.sigmoid(out.logits[:, 0])
+        is_novelty = self.sigmoid(out.logits[:, 1])
+
+        return ValNovOutput(
+            attentions=out.attentions,
+            hidden_states=out.hidden_states,
+            logits=out.logits,
+            loss=val_nov_loss(
+                is_val=is_validity,
+                is_nov=is_novelty,
+                should_val=should_validity,
+                should_nov=should_novelty,
+                weights=weights,
+                reduce=self.loss == "compute and reduce"
+            ) if self.loss != "ignore" and should_validity is not None and should_novelty is not None else None,
+            validity=is_validity,
+            novelty=is_novelty
+        )