<gh_stars>10-100
#
# Copyright (c) 2020 Saarland University.
#
# This file is part of AM Parser
# (see https://github.com/coli-saar/am-parser/).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from allennlp.nn.util import get_text_field_mask, get_range_vector
from allennlp.nn.util import get_device_of, get_lengths_from_binary_sequence_mask

from typing import Dict, Optional, Tuple, Any, List

import torch

from graph_dependency_parser.components.losses.base import EdgeLoss, EdgeExistenceLoss, EdgeLabelLoss


SMALL_FLOAT = 1e-13

def masked_log_softmax(vector: torch.Tensor, mask: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """
    Taken from AllenNLP version 0.8.4 and modified to use different definitions of "small number".
    On some machines the default of 1e-45 is too small: see https://github.com/coli-saar/am-parser/issues/94

    ``torch.nn.functional.log_softmax(vector)`` does not work if some elements of ``vector`` should be
    masked.  This performs a log_softmax on just the non-masked portions of ``vector``.  Passing
    ``None`` in for the mask is also acceptable; you'll just get a regular log_softmax.
    ``vector`` can have an arbitrary number of dimensions; the only requirement is that ``mask`` is
    broadcastable to ``vector's`` shape.  If ``mask`` has fewer dimensions than ``vector``, we will
    unsqueeze on dimension 1 until they match.  If you need a different unsqueezing of your mask,
    do it yourself before passing the mask into this function.
    In the case that the input vector is completely masked, the return value of this function is
    arbitrary, but not ``nan``.  You should be masking the result of whatever computation comes out
    of this in that case, anyway, so the specific values returned shouldn't matter.  Also, the way
    that we deal with this case relies on having single-precision floats; mixing half-precision
    floats with fully-masked vectors will likely give you ``nans``.
    If your logits are all extremely negative (i.e., the max value in your logit vector is -50 or
    lower), the way we handle masking here could mess you up.  But if you've got logit values that
    extreme, you've got bigger problems than this.
    """
    if mask is not None:
        mask = mask.float()
        while mask.dim() < vector.dim():
            mask = mask.unsqueeze(1)
        # vector + mask.log() is an easy way to zero out masked elements in logspace, but it
        # results in nans when the whole vector is masked.  We need a very small value instead of a
        # zero in the mask for these cases.  log(1 + 1e-45) is still basically 0, so we can safely
        # just add 1e-45 before calling mask.log().  We use 1e-45 because 1e-46 is so small it
        # becomes 0 - this is just the smallest value we can actually use.
        vector = vector + (mask + SMALL_FLOAT).log()
    return torch.nn.functional.log_softmax(vector, dim=dim)



@EdgeExistenceLoss.register("dm_edge_loss")
class DMLoss (EdgeExistenceLoss):
    """
    Dozat & Manning - Loss.
    """

    def loss(self, edge_scores: torch.Tensor,
                            head_indices: torch.Tensor,
                            mask: torch.Tensor) -> torch.Tensor:
        """
        Computes the edge loss for a sequence given gold head indices and tags.

        Parameters
        ----------
        edge_scores : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length, sequence_length) used to generate
            a distribution over attachments of a given word to all other words.
        head_indices : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length).
            The indices of the heads for every word.
        head_tags : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length).
            The dependency labels of the heads for every word.
        mask : ``torch.Tensor``, required.
            A mask of shape (batch_size, sequence_length), denoting unpadded
            elements in the sequence.

        Returns
        -------
        arc_nll : ``torch.Tensor``, required.
            The negative log likelihood from the arc loss.
        """
        float_mask = mask.float()
        batch_size, sequence_length, _ = edge_scores.size()
        # shape (batch_size, 1)
        range_vector = get_range_vector(batch_size, get_device_of(edge_scores)).unsqueeze(1)
        # shape (batch_size, sequence_length, sequence_length)
        normalised_arc_logits = masked_log_softmax(edge_scores,
                                                   mask) * float_mask.unsqueeze(2) * float_mask.unsqueeze(1)

        # index matrix with shape (batch, sequence_length)
        timestep_index = get_range_vector(sequence_length, get_device_of(edge_scores))
        child_index = timestep_index.view(1, sequence_length).expand(batch_size, sequence_length).long()
        # shape (batch_size, sequence_length)
        arc_loss = normalised_arc_logits[range_vector, child_index, head_indices]
        # We don't care about predictions for the symbolic ROOT token's head,
        # so we remove it from the loss.
        arc_loss = arc_loss[:, 1:]

        # The number of valid positions is equal to the number of unmasked elements minus
        # 1 per sequence in the batch, to account for the symbolic HEAD token.
        valid_positions = mask.sum() - batch_size

        arc_nll = -arc_loss.sum()
        if self.normalize_wrt_seq_len:
            arc_nll /= valid_positions.float()
        return arc_nll

@EdgeLabelLoss.register("dm_label_loss")
class DMLabelLoss(EdgeLabelLoss):

    def loss(self, edge_label_logits:torch.Tensor, mask:torch.Tensor, head_tags:torch.Tensor) -> torch.Tensor:
        """
        Computes the arc and tag loss for a sequence given gold head indices and tags.

        Parameters
        ----------
        edge_label_logits : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length, num_head_tags),
            that contains raw predictions for incoming edge labels
        head_tags : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length).
            The dependency labels of the heads for every word.
        mask : ``torch.Tensor``, required.
            A mask of shape (batch_size, sequence_length), denoting unpadded
            elements in the sequence.

        Returns
        -------
        tag_nll : ``torch.Tensor``, required.
            The negative log likelihood from the edge label loss.
        """
        float_mask = mask.float()
        batch_size, sequence_length, _ = edge_label_logits.size()
        # shape (batch_size, 1)
        range_vector = get_range_vector(batch_size, get_device_of(edge_label_logits)).unsqueeze(1)

        # shape (batch_size, sequence_length, num_head_tags)
        normalised_edge_label_logits = masked_log_softmax(edge_label_logits,
                                                        mask.unsqueeze(-1)) * float_mask.unsqueeze(-1)
        # index matrix with shape (batch, sequence_length)
        timestep_index = get_range_vector(sequence_length, get_device_of(edge_label_logits))
        child_index = timestep_index.view(1, sequence_length).expand(batch_size, sequence_length).long()
        # shape (batch_size, sequence_length)
        tag_loss = normalised_edge_label_logits[range_vector, child_index, head_tags]
        # We don't care about predictions for the symbolic ROOT token's head,
        # so we remove it from the loss.
        tag_loss = tag_loss[:, 1:]

        # The number of valid positions is equal to the number of unmasked elements minus
        # 1 per sequence in the batch, to account for the symbolic HEAD token.
        valid_positions = mask.sum() - batch_size

        if self.normalize_wrt_seq_len:
            return -tag_loss.sum()  / valid_positions.float()
        else:
            return -tag_loss.sum()

#
# Copyright (c) 2020 Saarland University.
#
# This file is part of AM Parser
# (see https://github.com/coli-saar/am-parser/).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from allennlp.nn.util import get_text_field_mask, get_range_vector
from allennlp.nn.util import get_device_of, get_lengths_from_binary_sequence_mask

from typing import Dict, Optional, Tuple, Any, List

import torch

from graph_dependency_parser.components.losses.base import EdgeLoss, EdgeExistenceLoss, EdgeLabelLoss


SMALL_FLOAT = 1e-13

def masked_log_softmax(vector: torch.Tensor, mask: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """
    Taken from AllenNLP version 0.8.4 and modified to use different definitions of "small number".
    On some machines the default of 1e-45 is too small: see https://github.com/coli-saar/am-parser/issues/94

    ``torch.nn.functional.log_softmax(vector)`` does not work if some elements of ``vector`` should be
    masked.  This performs a log_softmax on just the non-masked portions of ``vector``.  Passing
    ``None`` in for the mask is also acceptable; you'll just get a regular log_softmax.
    ``vector`` can have an arbitrary number of dimensions; the only requirement is that ``mask`` is
    broadcastable to ``vector's`` shape.  If ``mask`` has fewer dimensions than ``vector``, we will
    unsqueeze on dimension 1 until they match.  If you need a different unsqueezing of your mask,
    do it yourself before passing the mask into this function.
    In the case that the input vector is completely masked, the return value of this function is
    arbitrary, but not ``nan``.  You should be masking the result of whatever computation comes out
    of this in that case, anyway, so the specific values returned shouldn't matter.  Also, the way
    that we deal with this case relies on having single-precision floats; mixing half-precision
    floats with fully-masked vectors will likely give you ``nans``.
    If your logits are all extremely negative (i.e., the max value in your logit vector is -50 or
    lower), the way we handle masking here could mess you up.  But if you've got logit values that
    extreme, you've got bigger problems than this.
    """
    if mask is not None:
        mask = mask.float()
        while mask.dim() < vector.dim():
            mask = mask.unsqueeze(1)
        # vector + mask.log() is an easy way to zero out masked elements in logspace, but it
        # results in nans when the whole vector is masked.  We need a very small value instead of a
        # zero in the mask for these cases.  log(1 + 1e-45) is still basically 0, so we can safely
        # just add 1e-45 before calling mask.log().  We use 1e-45 because 1e-46 is so small it
        # becomes 0 - this is just the smallest value we can actually use.
        vector = vector + (mask + SMALL_FLOAT).log()
    return torch.nn.functional.log_softmax(vector, dim=dim)



@EdgeExistenceLoss.register("dm_edge_loss")
class DMLoss (EdgeExistenceLoss):
    """
    Dozat & Manning - Loss.
    """

    def loss(self, edge_scores: torch.Tensor,
                            head_indices: torch.Tensor,
                            mask: torch.Tensor) -> torch.Tensor:
        """
        Computes the edge loss for a sequence given gold head indices and tags.

        Parameters
        ----------
        edge_scores : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length, sequence_length) used to generate
            a distribution over attachments of a given word to all other words.
        head_indices : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length).
            The indices of the heads for every word.
        head_tags : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length).
            The dependency labels of the heads for every word.
        mask : ``torch.Tensor``, required.
            A mask of shape (batch_size, sequence_length), denoting unpadded
            elements in the sequence.

        Returns
        -------
        arc_nll : ``torch.Tensor``, required.
            The negative log likelihood from the arc loss.
        """
        float_mask = mask.float()
        batch_size, sequence_length, _ = edge_scores.size()
        # shape (batch_size, 1)
        range_vector = get_range_vector(batch_size, get_device_of(edge_scores)).unsqueeze(1)
        # shape (batch_size, sequence_length, sequence_length)
        normalised_arc_logits = masked_log_softmax(edge_scores,
                                                   mask) * float_mask.unsqueeze(2) * float_mask.unsqueeze(1)

        # index matrix with shape (batch, sequence_length)
        timestep_index = get_range_vector(sequence_length, get_device_of(edge_scores))
        child_index = timestep_index.view(1, sequence_length).expand(batch_size, sequence_length).long()
        # shape (batch_size, sequence_length)
        arc_loss = normalised_arc_logits[range_vector, child_index, head_indices]
        # We don't care about predictions for the symbolic ROOT token's head,
        # so we remove it from the loss.
        arc_loss = arc_loss[:, 1:]

        # The number of valid positions is equal to the number of unmasked elements minus
        # 1 per sequence in the batch, to account for the symbolic HEAD token.
        valid_positions = mask.sum() - batch_size

        arc_nll = -arc_loss.sum()
        if self.normalize_wrt_seq_len:
            arc_nll /= valid_positions.float()
        return arc_nll

@EdgeLabelLoss.register("dm_label_loss")
class DMLabelLoss(EdgeLabelLoss):

    def loss(self, edge_label_logits:torch.Tensor, mask:torch.Tensor, head_tags:torch.Tensor) -> torch.Tensor:
        """
        Computes the arc and tag loss for a sequence given gold head indices and tags.

        Parameters
        ----------
        edge_label_logits : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length, num_head_tags),
            that contains raw predictions for incoming edge labels
        head_tags : ``torch.Tensor``, required.
            A tensor of shape (batch_size, sequence_length).
            The dependency labels of the heads for every word.
        mask : ``torch.Tensor``, required.
            A mask of shape (batch_size, sequence_length), denoting unpadded
            elements in the sequence.

        Returns
        -------
        tag_nll : ``torch.Tensor``, required.
            The negative log likelihood from the edge label loss.
        """
        float_mask = mask.float()
        batch_size, sequence_length, _ = edge_label_logits.size()
        # shape (batch_size, 1)
        range_vector = get_range_vector(batch_size, get_device_of(edge_label_logits)).unsqueeze(1)

        # shape (batch_size, sequence_length, num_head_tags)
        normalised_edge_label_logits = masked_log_softmax(edge_label_logits,
                                                        mask.unsqueeze(-1)) * float_mask.unsqueeze(-1)
        # index matrix with shape (batch, sequence_length)
        timestep_index = get_range_vector(sequence_length, get_device_of(edge_label_logits))
        child_index = timestep_index.view(1, sequence_length).expand(batch_size, sequence_length).long()
        # shape (batch_size, sequence_length)
        tag_loss = normalised_edge_label_logits[range_vector, child_index, head_tags]
        # We don't care about predictions for the symbolic ROOT token's head,
        # so we remove it from the loss.
        tag_loss = tag_loss[:, 1:]

        # The number of valid positions is equal to the number of unmasked elements minus
        # 1 per sequence in the batch, to account for the symbolic HEAD token.
        valid_positions = mask.sum() - batch_size

        if self.normalize_wrt_seq_len:
            return -tag_loss.sum()  / valid_positions.float()
        else:
            return -tag_loss.sum()
# # Copyright (c) 2020 Saarland University. # # This file is part of AM Parser # (see https://github.com/coli-saar/am-parser/). # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Taken from AllenNLP version 0.8.4 and modified to use different definitions of "small number".     On some machines the default of 1e-45 is too small: see https://github.com/coli-saar/am-parser/issues/94      ``torch.nn.functional.log_softmax(vector)`` does not work if some elements of ``vector`` should be     masked.  This performs a log_softmax on just the non-masked portions of ``vector``.  Passing     ``None`` in for the mask is also acceptable; you'll just get a regular log_softmax.     ``vector`` can have an arbitrary number of dimensions; the only requirement is that ``mask`` is     broadcastable to ``vector's`` shape.  If ``mask`` has fewer dimensions than ``vector``, we will     unsqueeze on dimension 1 until they match.  If you need a different unsqueezing of your mask,     do it yourself before passing the mask into this function.     In the case that the input vector is completely masked, the return value of this function is     arbitrary, but not ``nan``.  You should be masking the result of whatever computation comes out     of this in that case, anyway, so the specific values returned shouldn't matter.  Also, the way     that we deal with this case relies on having single-precision floats; mixing half-precision     floats with fully-masked vectors will likely give you ``nans``.     If your logits are all extremely negative (i.e., the max value in your logit vector is -50 or     lower), the way we handle masking here could mess you up.  But if you've got logit values that     extreme, you've got bigger problems than this. # vector + mask.log() is an easy way to zero out masked elements in logspace, but it # results in nans when the whole vector is masked.  We need a very small value instead of a # zero in the mask for these cases.  log(1 + 1e-45) is still basically 0, so we can safely # just add 1e-45 before calling mask.log().  We use 1e-45 because 1e-46 is so small it # becomes 0 - this is just the smallest value we can actually use. Dozat & Manning - Loss. Computes the edge loss for a sequence given gold head indices and tags.          Parameters         ----------         edge_scores : ``torch.Tensor``, required.             A tensor of shape (batch_size, sequence_length, sequence_length) used to generate             a distribution over attachments of a given word to all other words.         head_indices : ``torch.Tensor``, required.             A tensor of shape (batch_size, sequence_length).             The indices of the heads for every word.         head_tags : ``torch.Tensor``, required.             A tensor of shape (batch_size, sequence_length).             The dependency labels of the heads for every word.         mask : ``torch.Tensor``, required.             A mask of shape (batch_size, sequence_length), denoting unpadded             elements in the sequence.          Returns         -------         arc_nll : ``torch.Tensor``, required.             The negative log likelihood from the arc loss. # shape (batch_size, 1) # shape (batch_size, sequence_length, sequence_length) # index matrix with shape (batch, sequence_length) # shape (batch_size, sequence_length) # We don't care about predictions for the symbolic ROOT token's head, # so we remove it from the loss. # The number of valid positions is equal to the number of unmasked elements minus # 1 per sequence in the batch, to account for the symbolic HEAD token. Computes the arc and tag loss for a sequence given gold head indices and tags.          Parameters         ----------         edge_label_logits : ``torch.Tensor``, required.             A tensor of shape (batch_size, sequence_length, num_head_tags),             that contains raw predictions for incoming edge labels         head_tags : ``torch.Tensor``, required.             A tensor of shape (batch_size, sequence_length).             The dependency labels of the heads for every word.         mask : ``torch.Tensor``, required.             A mask of shape (batch_size, sequence_length), denoting unpadded             elements in the sequence.          Returns         -------         tag_nll : ``torch.Tensor``, required.             The negative log likelihood from the edge label loss. # shape (batch_size, 1) # shape (batch_size, sequence_length, num_head_tags) # index matrix with shape (batch, sequence_length) # shape (batch_size, sequence_length) # We don't care about predictions for the symbolic ROOT token's head, # so we remove it from the loss. # The number of valid positions is equal to the number of unmasked elements minus # 1 per sequence in the batch, to account for the symbolic HEAD token.