Upload logic_consistency_loss.py with huggingface_hub

logic_consistency_loss.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LogicConsistencyLoss(nn.Module):
+    """Logic consistency loss.
+
+    Referenece:
+    "Leveraging Declarative Knowledge in Text and First-Order Logic for Fine-Grained Propaganda Detection".
+    https://arxiv.org/abs/2004.14201
+
+    Args:
+        use_sigmoid (bool): Can only be true for BCE based loss now.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'``: no reduction will be applied,
+            ``'mean'``: the sum of the output will be divided by the number of elements in the output,
+            ``'sum'``: the output will be summed.
+            Default: mean
+    """
+
+    def __init__(self, n_classes, reduce="max", reduction="mean"):
+        super(LogicConsistencyLoss, self).__init__()
+        if reduction not in ["sum", "mean", "none"]:
+            raise ValueError(
+                f"reduction should be 'sum', 'mean', or 'none', but got {reduction}"
+            )
+        if reduce not in ["max", "min", "avg"]:
+            raise ValueError(
+                f"reduce should be 'max', 'avg' or 'min', but got {reduce}"
+            )
+        self.n_classes = n_classes
+        self.reduce = reduce
+        self.reduction = reduction
+
+    def forward(self, fc, gc, gc_mask):
+        """Calculate the logic consistency loss.
+        Args:
+            fc (float tensor of size [batch_num, n_classes]):
+                Sentence level prediction with logits.
+            gc (float tensor of size [batch_num, token_num, n_classes]):
+                Token level prediction with logits.
+            gc_mask (float tensor of size [batch_num, token_num]):
+                Mask for token level prediction.
+        """
+        batch_num, token_num, class_num = gc.size()
+        assert (
+            class_num == self.n_classes
+        ), f"Class number mismatch: {class_num} vs {self.n_classes}"
+
+        fc = F.softmax(fc, dim=-1)[..., -1]  # [batch_num,]
+        gc = F.softmax(gc, dim=-1)[..., -1]  # [batch_num, token_num]
+        gc = torch.mul(gc, gc_mask)
+
+        if self.reduce == "max":
+            gc_maxpool, _ = torch.max(gc, dim=-1)  # [batch_num, ]
+        elif self.reduce == "min":
+            gc_maxpool, _ = torch.min(gc, dim=-1)
+        elif self.reduce == "avg":
+            gc_maxpool = torch.mean(gc, dim=-1)  # [batch_num, ]
+
+        pf = 1 - fc + torch.mul(fc, gc_maxpool)
+        pf = torch.clamp(pf, min=1e-8)  # Clamp to prevent log(0)
+
+        out = torch.neg(torch.log(pf))
+
+        if self.reduction == "mean":
+            loss = torch.mean(out)
+        elif self.reduction == "sum":
+            loss = torch.sum(out)
+        return loss