Source code for paddlenlp.layers.crf

# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

import paddle
import paddle.nn as nn

from ..utils.log import logger
from .sequence import sequence_mask

__all__ = ["LinearChainCrf", "LinearChainCrfLoss", "ViterbiDecoder"]


def log_sum_exp(vec, dim=0):
    # Avoid underflow and overflow
    max_num = paddle.max(vec, dim)
    max_exp = max_num.unsqueeze(-1)
    return max_num + paddle.log(paddle.sum(paddle.exp(vec - max_exp), dim))


[docs]class LinearChainCrf(nn.Layer): """ LinearChainCrf is a linear chain Conditional Random Field layer, it can implement sequential dependencies in the predictions. Therefore, it can take context into account whereas a classifier predicts a label for a single sample without considering "neighboring" samples. See https://repository.upenn.edu/cgi/viewcontent.cgi?article=1162&context=cis_papers for reference. Args: num_labels (int): The label number. crf_lr (float, optional): The crf layer learning rate. Defaults to ``0.1``. with_start_stop_tag (bool, optional): If set to True, the start tag and stop tag will be considered, the transitions params will be a tensor with a shape of `[num_labels+2, num_labels+2]`. Else, the transitions params will be a tensor with a shape of `[num_labels, num_labels]`. """ def __init__(self, num_labels, crf_lr=0.1, with_start_stop_tag=True): super(LinearChainCrf, self).__init__() if with_start_stop_tag: self.num_tags = num_labels + 2 # Additional [START] and [STOP] self.start_idx = int(self.num_tags - 1) self.stop_idx = int(self.num_tags - 2) else: self.num_tags = num_labels self.transitions = self.create_parameter( attr=paddle.ParamAttr(learning_rate=crf_lr), shape=[self.num_tags, self.num_tags], dtype="float32" ) self.with_start_stop_tag = with_start_stop_tag self._initial_alpha = None self._start_tensor = None self._stop_tensor = None self._batch_index = None self._seq_index = None self._batch_seq_index = None def _initialize_alpha(self, batch_size): # alpha accumulate the path value to get the different next tag if self._initial_alpha is None or batch_size > self._initial_alpha.shape[0]: # Initialized by a small value. initial_alpha = paddle.full((batch_size, self.num_tags - 1), dtype="float32", fill_value=-10000.0) # alpha_start fill_value = 0. > -10000., means the first one step START gets the most score. alpha_start = paddle.full((batch_size, 1), dtype="float32", fill_value=0.0) self._initial_alpha = paddle.concat([initial_alpha, alpha_start], axis=1) return self._initial_alpha[:batch_size, :]
[docs] def forward(self, inputs, lengths): """ Computes the normalization in a linear-chain CRF. See http://www.cs.columbia.edu/~mcollins/fb.pdf for reference. .. math:: F & = logZ(x) = log\\sum_y exp(score(x,y)) score(x,y) & = \\sum_i Emit(x_i,y_i) + Trans(y_{i-1}, y_i) p(y_i) & = Emit(x_i,y_i), T(y_{i-1}, y_i) = Trans(y_{i-1}, y_i) then we can get: .. math:: F(1) = log\\sum_{y1} exp(p(y_1) + T([START], y1)) .. math:: F(2) & = log\\sum_{y1}\\sum_{y2} exp(p(y_1) + T([START], y1) + p(y_2) + T(y_1,y_2)) \\\\ & = log\\sum_{y2} exp(F(1) + p(y_2) + T(y_1,y_2)) Further, We can get F(n) is a recursive formula with F(n-1). Args: inputs (Tensor): The input predicted tensor. Its dtype is float32 and has a shape of `[batch_size, sequence_length, num_tags]`. lengths (Tensor): The input length. Its dtype is int64 and has a shape of `[batch_size]`. Returns: Tensor: Returns the normalizers tensor `norm_score`. Its dtype is float32 and has a shape of `[batch_size]`. """ batch_size, seq_len, n_labels = inputs.shape inputs_t_exp = inputs.transpose([1, 0, 2]).unsqueeze(-1) # trans_exp: batch_size, num_tags, num_tags trans_exp = self.transitions.unsqueeze(0) all_alpha = [] if self.with_start_stop_tag: alpha = self._initialize_alpha(batch_size) for i, input_exp in enumerate(inputs_t_exp): # input_exp: batch_size, num_tags, num_tags # alpha_exp: batch_size, num_tags, num_tags if i == 0 and not self.with_start_stop_tag: alpha = inputs[:, 0] else: alpha_exp = alpha.unsqueeze(1) # F(n) = logsumexp(F(n-1) + p(y_n) + T(y_{n-1}, y_n)) mat = input_exp + trans_exp + alpha_exp alpha = log_sum_exp(mat, 2).squeeze(-1) all_alpha.append(alpha) # Get the valid alpha all_alpha = paddle.stack(all_alpha).transpose([1, 0, 2]) batch_index = self._get_batch_index(batch_size) last_index = lengths - 1 idxs = paddle.stack([batch_index, last_index], axis=1) alpha = paddle.gather_nd(all_alpha, idxs) if self.with_start_stop_tag: # The last one step alpha += self.transitions[self.stop_idx].unsqueeze(0) norm_score = log_sum_exp(alpha, 1) # .squeeze(-1) return norm_score
[docs] def gold_score(self, inputs, labels, lengths): """ Computes the unnormalized score for a tag sequence. $$ score(x,y) = \\sum_i Emit(x_i,y_i) + Trans(y_{i-1}, y_i) $$ Args: inputs (Tensor): The input predicted tensor. Its dtype is float32 and has a shape of `[batch_size, sequence_length, num_tags]`. labels (Tensor): The input label tensor. Its dtype is int64 and has a shape of `[batch_size, sequence_length]` lengths (Tensor): The input length. Its dtype is int64 and has a shape of `[batch_size]`. Returns: Tensor: Returns the unnormalized sequence scores tensor `unnorm_score`. Its dtype is float32 and has a shape of `[batch_size]`. """ unnorm_score = self._point_score(inputs, labels, lengths) + self._trans_score(labels, lengths) return unnorm_score
def _point_score(self, inputs, labels, lengths): batch_size, seq_len, n_labels = inputs.shape # Get the true label logit value flattened_inputs = inputs.reshape([-1]) offsets = paddle.unsqueeze(self._get_batch_index(batch_size) * seq_len * n_labels, 1) offsets += paddle.unsqueeze(self._get_seq_index(seq_len) * n_labels, 0) flattened_tag_indices = paddle.reshape(offsets + labels, [-1]) scores = paddle.gather(flattened_inputs, flattened_tag_indices).reshape([batch_size, seq_len]) mask = paddle.cast(sequence_mask(self._get_batch_seq_index(batch_size, seq_len), lengths), "float32") mask = mask[:, :seq_len] mask_scores = scores * mask score = paddle.sum(mask_scores, 1) return score def _trans_score(self, labels, lengths): batch_size, seq_len = labels.shape if self.with_start_stop_tag: # Add START and STOP on either side of the labels start_tensor, stop_tensor = self._get_start_stop_tensor(batch_size) labels_ext = paddle.concat([start_tensor, labels, stop_tensor], axis=1) mask = paddle.cast(sequence_mask(self._get_batch_seq_index(batch_size, seq_len), lengths + 1), "int64") pad_stop = paddle.full((batch_size, seq_len + 2), dtype="int64", fill_value=self.stop_idx) labels_ext = (1 - mask) * pad_stop + mask * labels_ext else: mask = paddle.cast(sequence_mask(self._get_batch_seq_index(batch_size, seq_len), lengths), "int64") labels_ext = labels start_tag_indices = labels_ext[:, :-1] stop_tag_indices = labels_ext[:, 1:] # Encode the indices in a flattened representation. transition_indices = start_tag_indices * self.num_tags + stop_tag_indices flattened_transition_indices = transition_indices.reshape([-1]) flattened_transition_params = paddle.flatten(self.transitions) scores = paddle.gather(flattened_transition_params, flattened_transition_indices).reshape([batch_size, -1]) mask_scores = scores * mask[:, 1:] # Accumulate the transition score score = paddle.sum(mask_scores, 1) return score def _get_start_stop_tensor(self, batch_size): if self._start_tensor is None or self._stop_tensor is None or batch_size != self._start_tensor.shape[0]: self._start_tensor = paddle.full((batch_size, 1), dtype="int64", fill_value=self.start_idx) self._stop_tensor = paddle.full((batch_size, 1), dtype="int64", fill_value=self.stop_idx) return self._start_tensor, self._stop_tensor def _get_batch_index(self, batch_size): if self._batch_index is None or batch_size != self._batch_index.shape[0]: self._batch_index = paddle.arange(end=batch_size, dtype="int64") return self._batch_index def _get_seq_index(self, length): if self._seq_index is None or length > self._seq_index.shape[0]: self._seq_index = paddle.arange(end=length, dtype="int64") return self._seq_index[:length] def _get_batch_seq_index(self, batch_size, length): if ( self._batch_seq_index is None or length + 2 > self._batch_seq_index.shape[1] or batch_size > self._batch_seq_index.shape[0] ): self._batch_seq_index = paddle.cumsum(paddle.ones([batch_size, length + 2], "int64"), axis=1) - 1 if self.with_start_stop_tag: return self._batch_seq_index[:batch_size, : length + 2] else: return self._batch_seq_index[:batch_size, :length]
[docs]class LinearChainCrfLoss(nn.Layer): """ The negative log-likelihood for linear chain Conditional Random Field (CRF). Args: crf (LinearChainCrf): The `LinearChainCrf` network object. Its parameter will be used to calculate the loss. """ def __init__(self, crf): super(LinearChainCrfLoss, self).__init__() self.crf = crf if isinstance(crf, paddle.Tensor): raise ValueError( "From paddlenlp >= 2.0.0b4, the first param of LinearChainCrfLoss shoule be a LinearChainCrf object. For input parameter 'crf.transitions', you can remove '.transitions' to 'crf'" )
[docs] def forward(self, inputs, lengths, labels, old_version_labels=None): """ Calculate the crf loss. Let $$ Z(x) = \\sum_{y'}exp(score(x,y')) $$, means the sum of all path scores, then we have $$ loss = -logp(y|x) = -log(exp(score(x,y))/Z(x)) = -score(x,y) + logZ(x) $$ Args: inputs (Tensor): The input predicted tensor. Its dtype is float32 and has a shape of `[batch_size, sequence_length, num_tags]`. lengths (Tensor): The input length. Its dtype is int64 and has a shape of `[batch_size]`. labels (Tensor) : The input label tensor. Its dtype is int64 and has a shape of `[batch_size, sequence_length]` old_version_labels (Tensor, optional): Unnecessary parameter for compatibility with older versions. Defaults to ``None``. Returns: Tensor: The crf loss. Its dtype is float32 and has a shape of `[batch_size]`. """ # Note: When closing to convergence, the loss could be a small negative number. This may caused by underflow when calculating exp in logsumexp. # We add relu here to avoid negative loss. In theory, the crf loss must be greater than or equal to 0, relu will not impact on it. if old_version_labels is not None: # TODO(qiujinxuan): rm compatibility support after lic. labels = old_version_labels if not getattr(self, "has_warn", False): logger.warning( "Compatibility Warning: The params of LinearChainCrfLoss.forward has been modified. The third param is `labels`, and the fourth is not necessary. Please update the usage." ) self.has_warn = True loss = nn.functional.relu(self.crf.forward(inputs, lengths) - self.crf.gold_score(inputs, labels, lengths)) return loss
[docs]class ViterbiDecoder(nn.Layer): """ ViterbiDecoder can decode the highest scoring sequence of tags, it should only be used at test time. Args: transitions (Tensor): The transition matrix. Its dtype is float32 and has a shape of `[num_tags, num_tags]`. with_start_stop_tag (bool, optional): If set to True, the last row and the last column of transitions will be considered as start tag, the penultimate row and the penultimate column of transitions will be considered as stop tag. Else, all the rows and columns will be considered as the real tag. Defaults to ``None``. """ def __init__(self, transitions, with_start_stop_tag=True): super(ViterbiDecoder, self).__init__() self.transitions = transitions self.with_start_stop_tag = with_start_stop_tag # If consider start and stop, -1 should be START and -2 should be STOP. if with_start_stop_tag: self.start_idx = -1 self.stop_idx = -2 self.num_tags = paddle.shape(transitions)[0] self._initial_alpha = None self._index = None self._batch_index = None self._batch_seq_index = None def _initialize_alpha(self, batch_size): # alpha accumulate the path value to get the different next tag if self._initial_alpha is None or batch_size > paddle.shape(self._initial_alpha)[0]: # Initialized by a small value. initial_alpha = paddle.full([batch_size, self.num_tags - 1], dtype="float32", fill_value=-10000.0) # alpha_start fill_value = 0. > -10000., means the first one step START gets the most score. alpha_start = paddle.full([batch_size, 1], dtype="float32", fill_value=0.0) self._initial_alpha = paddle.concat([initial_alpha, alpha_start], axis=1) return paddle.slice(self._initial_alpha, axes=[0], starts=[0], ends=[batch_size])
[docs] def forward(self, inputs, lengths): """ Decode the highest scoring sequence of tags. Args: inputs (Tensor): The unary emission tensor. Its dtype is float32 and has a shape of `[batch_size, sequence_length, num_tags]`. length (Tensor): The input length tensor storing real length of each sequence for correctness. Its dtype is int64 and has a shape of `[batch_size]`. Returns: tuple: Returns tuple (scores, paths). The `scores` tensor containing the score for the Viterbi sequence. Its dtype is float32 and has a shape of `[batch_size]`. The `paths` tensor containing the highest scoring tag indices. Its dtype is int64 and has a shape of `[batch_size, sequence_length]`. """ input_shape = paddle.shape(inputs) batch_size = input_shape[0] n_label = input_shape[2] inputs_t = inputs.transpose([1, 0, 2]) trans_exp = self.transitions.unsqueeze(0).expand([batch_size, n_label, n_label]) historys = [] left_length = lengths.clone() max_seq_len = left_length.max() # no need to expand the 'mask' in the following iteration left_length = left_length.unsqueeze(-1).expand([batch_size, n_label]) if self.with_start_stop_tag: alpha = self._initialize_alpha(batch_size) else: alpha = paddle.zeros((batch_size, self.num_tags), dtype="float32") for i, logit in enumerate(inputs_t[:max_seq_len]): # if not with_start_stop_tag, the first label has not antecedent tag. if i == 0 and not self.with_start_stop_tag: alpha = logit left_length = left_length - 1 continue alpha_exp = alpha.unsqueeze(2) # alpha_trn_sum: batch_size, n_labels, n_labels alpha_trn_sum = alpha_exp + trans_exp # alpha_max: batch_size, n_labels # We don't include the emission scores here because the max does not depend on them (we add them in below) alpha_max = alpha_trn_sum.max(1) # If with_start_stop_tag, the first antecedent tag must be START, else the first label has not antecedent tag. # So we can record the path from i=1. if i >= 1: alpha_argmax = alpha_trn_sum.argmax(1) historys.append(alpha_argmax) # Now add the emission scores alpha_nxt = alpha_max + logit mask = paddle.cast((left_length > 0), dtype="float32") alpha = mask * alpha_nxt + (1 - mask) * alpha if self.with_start_stop_tag: mask = paddle.cast((left_length == 1), dtype="float32") alpha += mask * trans_exp[:, self.stop_idx] left_length = left_length - 1 # last_ids: batch_size scores, last_ids = alpha.max(1), alpha.argmax(1) if max_seq_len == 1: return scores, last_ids.unsqueeze(1) # Trace back the best path # historys: seq_len, batch_size, n_labels historys = paddle.stack(historys) left_length = left_length[:, 0] tag_mask = paddle.cast((left_length >= 0), "int64") last_ids_update = last_ids * tag_mask batch_path = [last_ids_update] batch_offset = self._get_batch_index(batch_size) * n_label historys = paddle.reverse(historys, [0]) for hist in historys: # hist: batch_size, n_labels left_length = left_length + 1 gather_idx = batch_offset + last_ids tag_mask = paddle.cast((left_length > 0), "int64") last_ids_update = paddle.gather(hist.flatten(), gather_idx) * tag_mask zero_len_mask = paddle.cast((left_length == 0), "int64") last_ids_update = last_ids_update * (1 - zero_len_mask) + last_ids * zero_len_mask batch_path.append(last_ids_update) tag_mask = paddle.cast((left_length >= 0), "int64") last_ids = last_ids_update + last_ids * (1 - tag_mask) batch_path = paddle.reverse(paddle.stack(batch_path, 1), [1]) return scores, batch_path
def _get_batch_index(self, batch_size): if self._batch_index is None or batch_size != paddle.shape(self._batch_index)[0]: self._batch_index = paddle.arange(end=batch_size, dtype="int64") return self._batch_index