def feed_batch(self, h: torch.FloatTensor, batch: Dict[str, torch.Tensor],
mask: torch.BoolTensor, decoder: torch.nn.Module):
mask = self.compute_mask(batch)
output_dict = decoder(h, batch, mask)
if decoder.training:
mask = mask.clone()
mask[:, 0] = 0
return output_dict, mask
def feed_batch(self, h: torch.FloatTensor, batch: Dict[str, torch.Tensor], mask: torch.BoolTensor,
decoder: torch.nn.Module):
logits = super().feed_batch(h, batch, mask, decoder)
arc_scores = logits[0]
mask = mask.clone()
mask[:, 0] = 0
mask = self.convert_to_3d_mask(arc_scores, mask)
punct_mask = self.convert_to_3d_puncts(batch.get('punct_mask', None), mask)
return logits, mask, punct_mask
def feed_batch(self, h: torch.FloatTensor, batch: Dict[str, torch.Tensor], mask: torch.BoolTensor,
decoder: torch.nn.Module):
logits = super().feed_batch(h, batch, mask, decoder)
mask = mask.clone()
mask[:, 0] = 0
return logits, mask
