def compute_loss(self, batch: Dict[str, Any],
output: Union[torch.Tensor, Dict[str, torch.Tensor], Iterable[torch.Tensor], Any], criterion) -> \
Union[torch.FloatTensor, Dict[str, torch.FloatTensor]]:
(arc_scores, rel_scores), mask = output
return BiaffineDependencyParser.compute_loss(self, arc_scores, rel_scores, batch['arc'], batch['rel_id'], mask,
criterion,
batch)
def compute_loss(self, arc_scores, rel_scores, arcs, rels, mask, criterion, batch=None):
parse_loss = BiaffineDependencyParser.compute_loss(self, arc_scores, rel_scores, arcs, rels, mask, criterion, batch)
if self.model.training:
gold_input_ids = batch['gold_input_ids']
pred_input_ids = batch['pred_input_ids']
input_ids_mask = batch['input_ids_mask']
token_span = batch['token_span']
gold_input_ids = batch['gold_input_ids'] = gold_input_ids.gather(1, token_span[:, :, 0])
input_ids_mask = batch['input_ids_mask'] = input_ids_mask.gather(1, token_span[:, :, 0])
mlm_loss = F.cross_entropy(pred_input_ids[input_ids_mask], gold_input_ids[input_ids_mask])
loss = parse_loss + mlm_loss
return loss
return parse_loss
def forward(self,
hidden,
batch: Dict[str, torch.Tensor],
mask) -> Dict[str, Any]:
mask_without_root = mask.clone()
mask_without_root[:, 0] = False
logits = {}
class_probabilities = {}
output_dict = {"logits": logits,
"class_probabilities": class_probabilities}
loss = 0
arc = batch.get('arc', None)
# Run through each of the tasks on the shared encoder and save predictions
for task in self.decoders:
if self.scalar_mix:
decoder_input = self.scalar_mix[task](hidden, mask)
else:
decoder_input = hidden
if task == "deps":
s_arc, s_rel = self.decoders[task](decoder_input, mask)
pred_output = {'class_probabilities': {'s_arc': s_arc, 's_rel': s_rel}}
if arc is not None:
# noinspection PyTypeChecker
pred_output['loss'] = BiaffineDependencyParser.compute_loss(None, s_arc, s_rel, arc,
batch['rel_id'],
mask_without_root,
torch.nn.functional.cross_entropy)
else:
pred_output = self.decoders[task](decoder_input, mask_without_root,
batch.get(self.gold_keys[task], None))
if 'logits' in pred_output:
logits[task] = pred_output["logits"]
if 'class_probabilities' in pred_output:
class_probabilities[task] = pred_output["class_probabilities"]
if 'loss' in pred_output:
# Keep track of the loss if we have the gold tags available
loss += pred_output["loss"]
if arc is not None:
output_dict["loss"] = loss
return output_dict