Fix bug of incorrect masking for ConstituencyCRF

yzhangcs · yzhangcs · commit aa83656eed02 · 2021-11-05T03:00:39.000+08:00
diff --git a/supar/models/const.py b/supar/models/const.py
@@ -190,7 +190,7 @@ def loss(self, s_span, s_label, charts, mask, mbr=True):
         """
 
         span_mask = charts.ge(0) & mask
-        span_dist = ConstituencyCRF(s_span, mask[:, 0].sum())
+        span_dist = ConstituencyCRF(s_span, mask[:, 0].sum(-1))
         span_loss = -span_dist.log_prob(span_mask).sum() / mask[:, 0].sum()
         span_probs = span_dist.marginals if mbr else s_span
         label_loss = self.criterion(s_label[span_mask], charts[span_mask])
@@ -213,7 +213,7 @@ def decode(self, s_span, s_label, mask):
                 Sequences of factorized labeled trees traversed in pre-order.
         """
 
-        span_preds = ConstituencyCRF(s_span, mask[:, 0].sum()).argmax
+        span_preds = ConstituencyCRF(s_span, mask[:, 0].sum(-1)).argmax
         label_preds = s_label.argmax(-1).tolist()
         return [[(i, j, labels[i][j]) for i, j in spans] for spans, labels in zip(span_preds, label_preds)]
 
@@ -439,6 +439,6 @@ def decode(self, s_span, s_label, mask):
                 Sequences of factorized labeled trees traversed in pre-order.
         """
 
-        span_preds = ConstituencyCRF(s_span, mask[:, 0].sum()).argmax
+        span_preds = ConstituencyCRF(s_span, mask[:, 0].sum(-1)).argmax
         label_preds = s_label.argmax(-1).tolist()
         return [[(i, j, labels[i][j]) for i, j in spans] for spans, labels in zip(span_preds, label_preds)]
diff --git a/supar/parsers/const.py b/supar/parsers/const.py
@@ -213,7 +213,7 @@ def _predict(self, loader):
             mask = (mask.unsqueeze(1) & mask.unsqueeze(2)).triu_(1)
             lens = mask[:, 0].sum(-1)
             s_span, s_label = self.model(words, feats)
-            s_span = ConstituencyCRF(s_span, mask[:, 0].sum()).marginals if self.args.mbr else s_span
+            s_span = ConstituencyCRF(s_span, mask[:, 0].sum(-1)).marginals if self.args.mbr else s_span
             chart_preds = self.model.decode(s_span, s_label, mask)
             preds['trees'].extend([Tree.build(tree, [(i, j, self.CHART.vocab[label]) for i, j, label in chart])
                                    for tree, chart in zip(trees, chart_preds)])
