def decode(self,
s_arc,
s_sib,
s_rel,
mask,
tree=False,
mbr=True,
proj=False):
"""
Args:
s_arc (torch.Tensor): [batch_size, seq_len, seq_len]
The scores of all possible arcs.
s_sib (torch.Tensor): [batch_size, seq_len, seq_len, seq_len]
The scores of all possible dependent-head-sibling triples.
s_rel (torch.Tensor): [batch_size, seq_len, seq_len, n_labels]
The scores of all possible labels on each arc.
mask (torch.BoolTensor): [batch_size, seq_len, seq_len]
Mask for covering the unpadded tokens.
tree (bool):
If True, ensures to output well-formed trees. Default: False.
mbr (bool):
If True, performs MBR decoding. Default: True.
proj (bool):
If True, ensures to output projective trees. Default: False.
Returns:
arc_preds (torch.Tensor): [batch_size, seq_len]
The predicted arcs.
rel_preds (torch.Tensor): [batch_size, seq_len]
The predicted labels.
"""
lens = mask.sum(1)
# prevent self-loops
s_arc.diagonal(0, 1, 2).fill_(float('-inf'))
arc_preds = s_arc.argmax(-1)
bad = [
not CoNLL.istree(seq[1:i + 1], proj)
for i, seq in zip(lens.tolist(), arc_preds.tolist())
]
if tree and any(bad):
if proj and not mbr:
arc_preds = eisner2o((s_arc, s_sib), mask)
else:
alg = eisner if proj else mst
arc_preds[bad] = alg(s_arc[bad], mask[bad])
rel_preds = s_rel.argmax(-1).gather(
-1, arc_preds.unsqueeze(-1)).squeeze(-1)
return arc_preds, rel_preds
def decode(self,
s_arc,
s_sib,
s_rel,
mask,
tree=False,
mbr=True,
proj=False):
r"""
Args:
s_arc (~torch.Tensor): ``[batch_size, seq_len, seq_len]``.
Scores of all possible arcs.
s_sib (~torch.Tensor): ``[batch_size, seq_len, seq_len, seq_len]``.
Scores of all possible dependent-head-sibling triples.
s_rel (~torch.Tensor): ``[batch_size, seq_len, seq_len, n_labels]``.
Scores of all possible labels on each arc.
mask (~torch.BoolTensor): ``[batch_size, seq_len]``.
The mask for covering the unpadded tokens.
tree (bool):
If ``True``, ensures to output well-formed trees. Default: ``False``.
mbr (bool):
If ``True``, performs MBR decoding. Default: ``True``.
proj (bool):
If ``True``, ensures to output projective trees. Default: ``False``.
Returns:
~torch.LongTensor, ~torch.LongTensor:
Predicted arcs and labels of shape ``[batch_size, seq_len]``.
"""
lens = mask.sum(1)
arc_preds = s_arc.argmax(-1)
bad = [
not CoNLL.istree(seq[1:i + 1], proj)
for i, seq in zip(lens.tolist(), arc_preds.tolist())
]
if tree and any(bad):
if proj and not mbr:
arc_preds = eisner2o((s_arc, s_sib), mask)
else:
alg = eisner if proj else mst
arc_preds[bad] = alg(s_arc[bad], mask[bad])
rel_preds = s_rel.argmax(-1).gather(
-1, arc_preds.unsqueeze(-1)).squeeze(-1)
return arc_preds, rel_preds