add "loss_attention"

Author: sfluegel

chebai/loss/semantic.py

@@ -18,6 +18,7 @@ def __init__(
         impl_loss_weight=0.1,  # weight of implication loss in relation to base_loss
         pos_scalar=1,
         pos_epsilon=0.01,
+        loss_attention=False,
     ):
         super().__init__()
         self.data_extractor = data_extractor
@@ -36,6 +37,7 @@ def __init__(
         self.impl_weight = impl_loss_weight
         self.pos_scalar = pos_scalar
         self.eps = pos_epsilon
+        self.loss_attention = loss_attention
 
     def forward(self, input, target, **kwargs):
         nnl = kwargs.pop("non_null_labels", None)
@@ -82,6 +84,8 @@ def _calculate_implication_loss(self, l, r):
         else:
             raise NotImplementedError(f"Unknown tnorm {self.tnorm}")
 
+        if self.loss_attention:
+            individual_loss *= individual_loss.softmax(dim=-1)
         return torch.mean(
             torch.sum(individual_loss, dim=-1),
             dim=0,