wehub-resource-sync
1248 lines
54 KiB
Python
1248 lines
54 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. team.
|
|
#
|
|
# 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 math
|
|
from typing import Optional, Tuple
|
|
|
|
import paddle
|
|
import paddle.nn as nn
|
|
import paddle.nn.functional as F
|
|
from paddle import Tensor
|
|
from paddle.nn import Layer
|
|
|
|
from .. import PretrainedModel, register_base_model
|
|
from ..activations import ACT2FN
|
|
from .configuration import (
|
|
PROPHETNET_PRETRAINED_INIT_CONFIGURATION,
|
|
PROPHETNET_PRETRAINED_RESOURCE_FILES_MAP,
|
|
ProphetNetConfig,
|
|
)
|
|
|
|
__all__ = [
|
|
"ProphetNetModel",
|
|
"ProphetNetPretrainedModel",
|
|
"ProphetNetEncoder",
|
|
"ProphetNetDecoder",
|
|
"ProphetNetForConditionalGeneration",
|
|
]
|
|
|
|
|
|
def ngram_attention_bias(sequence_length, ngram, dtype):
|
|
"""
|
|
This function computes the bias for the predict stream
|
|
"""
|
|
left_block = paddle.ones((ngram, sequence_length, sequence_length), dtype=dtype) * float("-inf")
|
|
right_block = left_block.detach().clone()
|
|
# create bias
|
|
for stream_idx in range(ngram):
|
|
right_block[stream_idx] = right_block[stream_idx].fill_diagonal_(0, wrap=False)
|
|
left_block[stream_idx] = paddle.triu(left_block[stream_idx], diagonal=-stream_idx + 1)
|
|
|
|
left_block[:, :, 0] = 0
|
|
return paddle.concat([left_block, right_block], axis=2)
|
|
|
|
|
|
def compute_relative_buckets(num_buckets, max_distance, relative_positions, is_bidirectional=False):
|
|
"""
|
|
This function computes individual parts of the relative position buckets. For more detail, see paper.
|
|
"""
|
|
inv_relative_positions = -relative_positions
|
|
rel_positions_bucket = 0
|
|
|
|
if is_bidirectional:
|
|
num_buckets = num_buckets // 2
|
|
rel_positions_bucket = (
|
|
rel_positions_bucket
|
|
+ paddle.cast(
|
|
paddle.less_than(inv_relative_positions, paddle.zeros_like(inv_relative_positions)), dtype=paddle.int32
|
|
)
|
|
* num_buckets
|
|
)
|
|
inv_relative_positions = paddle.abs(inv_relative_positions)
|
|
else:
|
|
inv_relative_positions = paddle.cast(
|
|
paddle.less_than(paddle.zeros_like(inv_relative_positions), inv_relative_positions), dtype=paddle.int32
|
|
) * inv_relative_positions.astype(paddle.int32)
|
|
|
|
max_exact = num_buckets // 2
|
|
is_small = paddle.less_than(inv_relative_positions, paddle.to_tensor(max_exact).cast(dtype=paddle.int32))
|
|
val_if_large = max_exact + paddle.log(
|
|
paddle.cast(inv_relative_positions, dtype=paddle.float32) / max_exact
|
|
) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
|
|
val_if_large_num_buckets = paddle.ones_like(val_if_large) * (num_buckets - 1)
|
|
val_if_large_lt = paddle.cast(paddle.less_than(val_if_large, val_if_large_num_buckets), dtype=paddle.int32)
|
|
val_if_large = val_if_large_lt * val_if_large.astype(val_if_large_lt.dtype) + (
|
|
1 - val_if_large_lt
|
|
) * val_if_large_num_buckets.astype(val_if_large_lt.dtype)
|
|
rel_positions_bucket = rel_positions_bucket + paddle.where(
|
|
is_small, paddle.cast(inv_relative_positions, dtype=paddle.int32), val_if_large
|
|
)
|
|
return rel_positions_bucket
|
|
|
|
|
|
def compute_all_stream_relative_buckets(num_buckets, max_distance, position_ids):
|
|
"""
|
|
This function computes both main and predict relative position buckets. For more detail, see paper.
|
|
"""
|
|
# main stream
|
|
main_stream_relative_positions = paddle.tile(
|
|
paddle.unsqueeze(position_ids, axis=1), repeat_times=[1, position_ids.shape[-1], 1]
|
|
)
|
|
main_stream_relative_positions = main_stream_relative_positions - paddle.unsqueeze(position_ids, axis=-1)
|
|
|
|
# predicting stream
|
|
predicting_stream_relative_positions = paddle.unsqueeze(
|
|
paddle.concat([position_ids - 1, position_ids], axis=-1), axis=1
|
|
)
|
|
predicting_stream_relative_positions = paddle.tile(
|
|
predicting_stream_relative_positions, repeat_times=[1, position_ids.shape[-1], 1]
|
|
)
|
|
predicting_stream_relative_positions = predicting_stream_relative_positions - paddle.unsqueeze(
|
|
position_ids, axis=-1
|
|
)
|
|
|
|
# get both position buckets
|
|
main_relative_position_buckets = compute_relative_buckets(
|
|
num_buckets, max_distance, main_stream_relative_positions, is_bidirectional=False
|
|
)
|
|
predict_relative_position_buckets = compute_relative_buckets(
|
|
num_buckets, max_distance, predicting_stream_relative_positions, is_bidirectional=False
|
|
)
|
|
return main_relative_position_buckets, predict_relative_position_buckets
|
|
|
|
|
|
class ProphetNetPretrainedModel(PretrainedModel):
|
|
"""
|
|
An abstract class for pretrained Prophetnet models. It provides Prophetnet related
|
|
`model_config_file`, `pretrained_init_configuration`, `resource_files_names`,
|
|
`pretrained_resource_files_map`, `base_model_prefix` for downloading and
|
|
loading pretrained models.
|
|
"""
|
|
|
|
pretrained_init_configuration = PROPHETNET_PRETRAINED_INIT_CONFIGURATION
|
|
pretrained_resource_files_map = PROPHETNET_PRETRAINED_RESOURCE_FILES_MAP
|
|
base_model_prefix = "prophetnet"
|
|
config_class = ProphetNetConfig
|
|
|
|
def _init_weights(self, layer):
|
|
if isinstance(layer, nn.Linear):
|
|
layer.weight.set_value(
|
|
paddle.tensor.normal(
|
|
mean=0.0,
|
|
std=self.config.init_std,
|
|
shape=layer.weight.shape,
|
|
)
|
|
)
|
|
if layer.bias is not None:
|
|
layer.bias.set_value(paddle.tensor.zeros(layer.bias.shape))
|
|
|
|
def _shift_right(self, input_ids):
|
|
decoder_start_token_id = self.config.decoder_start_token_id
|
|
pad_token_id = self.config.pad_token_id
|
|
|
|
assert decoder_start_token_id is not None, (
|
|
"self.config.decoder_start_token_id has to be defined. "
|
|
"In ProphetNet it is usually set to the pad_token_id. See ProphetNet docs for more information"
|
|
)
|
|
|
|
# shift inputs to the right
|
|
shifted_input_ids = paddle.zeros_like(input_ids)
|
|
shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
|
|
shifted_input_ids[..., 0] = decoder_start_token_id
|
|
|
|
assert pad_token_id is not None, "self.config.pad_token_id has to be defined."
|
|
# replace possible -100 values in labels by `pad_token_id`
|
|
shifted_input_ids_mask = paddle.cast(shifted_input_ids == -100, dtype=paddle.int32)
|
|
shifted_input_ids = shifted_input_ids_mask * pad_token_id + (1 - shifted_input_ids_mask) * shifted_input_ids
|
|
|
|
assert (
|
|
paddle.sum(paddle.cast(shifted_input_ids >= 0, dtype=paddle.int32)).item() == shifted_input_ids.shape[-1]
|
|
), "Verify that `shifted_input_ids` has only positive values"
|
|
|
|
return shifted_input_ids
|
|
|
|
|
|
class ProphetNetPositionalEmbeddings(nn.Embedding):
|
|
"""
|
|
ProphetNetPositional Embeddings.
|
|
"""
|
|
|
|
def __init__(self, config: ProphetNetConfig):
|
|
self.max_length = config.max_position_embeddings
|
|
super(ProphetNetPositionalEmbeddings, self).__init__(
|
|
config.max_position_embeddings, config.hidden_size, config.pad_token_id
|
|
)
|
|
|
|
def forward(self, inputs_shape, attention_mask=None, past_key_values=None, position_ids=None):
|
|
assert (position_ids is None) or (
|
|
self._padding_idx is None
|
|
), "If position_ids is pre-computed then padding_idx should not be set."
|
|
|
|
if position_ids is None:
|
|
if past_key_values is not None:
|
|
# position_ids is the same for every token when decoding a single step
|
|
# Without the int() cast, it doesn't work in some cases when exporting to ONNX
|
|
prev_num_input_ids = past_key_values[0][0].shape[2]
|
|
num_input_ids = inputs_shape[1] + prev_num_input_ids
|
|
position_ids = paddle.ones((1, 1), dtype="int64") * (int(self._padding_idx + num_input_ids))
|
|
else:
|
|
if attention_mask is None:
|
|
attention_mask = paddle.ones(inputs_shape, dtype="int64")
|
|
|
|
# retrieve position_ids from input_ids / attention_mask
|
|
position_ids = (
|
|
paddle.cast(
|
|
paddle.cast(paddle.cumsum(attention_mask, axis=1), dtype=attention_mask.dtype)
|
|
* attention_mask,
|
|
dtype=paddle.int64,
|
|
)
|
|
+ self._padding_idx
|
|
)
|
|
|
|
# make sure position_ids are not bigger then max_length
|
|
position_ids = paddle.clip(position_ids, min=0, max=self.max_length - 1)
|
|
|
|
return super().forward(position_ids), position_ids
|
|
|
|
def _forward(self, position_ids):
|
|
return super().forward(position_ids)
|
|
|
|
|
|
class ProphetNetAttention(Layer):
|
|
"""
|
|
Multi-headed attention from 'Attention Is All You Need' paper.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, attention_dropout, dropout, num_attn_heads: int):
|
|
super().__init__()
|
|
hidden_size = hidden_size
|
|
|
|
self.attention_dropout = attention_dropout
|
|
self.dropout = dropout
|
|
self.num_attn_heads = num_attn_heads
|
|
self.head_dim = hidden_size // num_attn_heads
|
|
|
|
assert (
|
|
self.head_dim * num_attn_heads == hidden_size
|
|
), "`config.hidden_size` must be divisible by `config.num_encoder_attention_heads` and `config.num_decoder_attention_heads`"
|
|
|
|
self.key_proj = nn.Linear(hidden_size, hidden_size)
|
|
self.value_proj = nn.Linear(hidden_size, hidden_size)
|
|
self.query_proj = nn.Linear(hidden_size, hidden_size)
|
|
|
|
self.out_proj = nn.Linear(hidden_size, hidden_size)
|
|
|
|
def _shape(self, tensor: paddle.Tensor, seq_len: int, bsz: int):
|
|
return paddle.transpose(
|
|
paddle.reshape(tensor, [bsz, seq_len, self.num_attn_heads, self.head_dim]), (0, 2, 1, 3)
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
key_value_states: Optional[Tensor] = None,
|
|
attention_mask: Optional[Tensor] = None,
|
|
past_key_value: Optional[Tuple[Tensor]] = None,
|
|
) -> Tuple[Tensor, Optional[Tensor]]:
|
|
|
|
batch_size, tgt_len, hidden_size = hidden_states.shape
|
|
|
|
# if key_value_states are provided this layer is used as a cross-attention layer
|
|
# for the decoder
|
|
is_cross_attention = key_value_states is not None
|
|
assert hidden_states.shape == [
|
|
batch_size,
|
|
tgt_len,
|
|
hidden_size,
|
|
], f"Size of hidden states should be {batch_size, tgt_len, hidden_size}, but is {hidden_states.shape}"
|
|
|
|
# previous time steps are cached - no need to recompute key and value if they are static
|
|
query_states = self.query_proj(hidden_states) / (self.head_dim**0.5)
|
|
|
|
if is_cross_attention and past_key_value is not None:
|
|
# reuse k,v, cross_attentions
|
|
key_states = past_key_value[0]
|
|
value_states = past_key_value[1]
|
|
elif is_cross_attention:
|
|
# cross_attentions
|
|
key_states = self._shape(self.key_proj(key_value_states), -1, batch_size)
|
|
value_states = self._shape(self.value_proj(key_value_states), -1, batch_size)
|
|
else:
|
|
# self_attention
|
|
key_states = self._shape(self.key_proj(hidden_states), -1, batch_size)
|
|
value_states = self._shape(self.value_proj(hidden_states), -1, batch_size)
|
|
|
|
if is_cross_attention:
|
|
# if cross_attention save Tuple(paddle.Tensor, paddle.Tensor) of all cross attention key/value_states.
|
|
# Further calls to cross_attention layer can then reuse all cross-attention
|
|
# key/value_states (first "if" case)
|
|
# if encoder bi-directional self-attention `past_key_value` is always `None`
|
|
past_key_value = (key_states, value_states)
|
|
|
|
# project states into the correct shape
|
|
proj_shape = (batch_size * self.num_attn_heads, -1, self.head_dim)
|
|
query_states = paddle.reshape(self._shape(query_states, tgt_len, batch_size), proj_shape)
|
|
key_states = paddle.reshape(key_states, proj_shape)
|
|
value_states = paddle.reshape(value_states, proj_shape)
|
|
|
|
src_len = key_states.shape[1]
|
|
attn_weights = paddle.bmm(query_states, key_states.transpose((0, 2, 1)))
|
|
assert attn_weights.shape == [
|
|
batch_size * self.num_attn_heads,
|
|
tgt_len,
|
|
src_len,
|
|
], f"`attn_weights` should be of size {batch_size * self.num_attn_heads, tgt_len, src_len}, but is of size {attn_weights.shape}"
|
|
|
|
# This is part of a workaround to get around fork/join parallelism not supporting Optional types.
|
|
if attention_mask is not None and len(attention_mask.shape) == 0:
|
|
attention_mask = None
|
|
assert attention_mask is None or attention_mask.shape == [
|
|
self.num_attn_heads * batch_size,
|
|
1,
|
|
src_len,
|
|
], f"`attention_mask` should be `None` or of shape attention_mask.shape == {batch_size * self.num_attn_heads, 1, src_len}, but is {attention_mask.shape}"
|
|
|
|
if attention_mask is not None: # don't attend to padding symbols
|
|
attn_weights = attn_weights + attention_mask
|
|
|
|
attn_weights = F.softmax(attn_weights, axis=-1)
|
|
|
|
attn_probs = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
|
|
|
|
attn_output = paddle.bmm(attn_probs, value_states)
|
|
assert attn_output.shape == [
|
|
batch_size * self.num_attn_heads,
|
|
tgt_len,
|
|
self.head_dim,
|
|
], f"`attn_output` should be of shape {batch_size * self.num_attn_heads, tgt_len, self.head_dim}, but is of shape {attn_output.shape}"
|
|
|
|
attn_output = paddle.reshape(
|
|
paddle.transpose(
|
|
paddle.reshape(attn_output, (batch_size, self.num_attn_heads, tgt_len, self.head_dim)), (0, 2, 1, 3)
|
|
),
|
|
(batch_size, tgt_len, hidden_size),
|
|
)
|
|
|
|
attn_output = self.out_proj(attn_output)
|
|
|
|
attn_output = F.dropout(attn_output, p=self.dropout, training=self.training)
|
|
return attn_output, past_key_value
|
|
|
|
|
|
class ProphetNetFeedForward(Layer):
|
|
"""
|
|
This is the residual two feed-forward layer block based on the original Transformer implementation.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, activation_function, activation_dropout, dropout, ffn_dim: int):
|
|
super(ProphetNetFeedForward, self).__init__()
|
|
self.activation_fn = ACT2FN[activation_function]
|
|
self.intermediate = nn.Linear(hidden_size, ffn_dim)
|
|
self.output = nn.Linear(ffn_dim, hidden_size)
|
|
self.activation_dropout = activation_dropout
|
|
self.dropout = dropout
|
|
|
|
def forward(self, hidden_states):
|
|
hidden_states = self.intermediate(hidden_states)
|
|
hidden_states = self.activation_fn(hidden_states)
|
|
|
|
hidden_states = F.dropout(hidden_states, p=self.activation_dropout, training=self.training)
|
|
hidden_states = self.output(hidden_states)
|
|
hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
return hidden_states
|
|
|
|
|
|
class ProphetNetNgramSelfAttention(Layer):
|
|
def __init__(
|
|
self,
|
|
hidden_size,
|
|
num_buckets,
|
|
relative_max_distance,
|
|
num_decoder_attention_heads,
|
|
dropout,
|
|
attention_dropout,
|
|
ngram,
|
|
):
|
|
super(ProphetNetNgramSelfAttention, self).__init__()
|
|
|
|
self.hidden_size = hidden_size
|
|
|
|
self.num_buckets = num_buckets
|
|
self.relative_max_distance = relative_max_distance
|
|
self.num_attn_heads = num_decoder_attention_heads
|
|
self.dropout = dropout
|
|
self.attention_dropout = attention_dropout
|
|
self.head_dim = hidden_size // self.num_attn_heads
|
|
self.ngram = ngram
|
|
|
|
assert (
|
|
self.head_dim * self.num_attn_heads == hidden_size
|
|
), "config.hidden_size must be divisible by num_attn_heads"
|
|
# key, value, query projection
|
|
self.key_proj = nn.Linear(hidden_size, hidden_size)
|
|
self.value_proj = nn.Linear(hidden_size, hidden_size)
|
|
self.query_proj = nn.Linear(hidden_size, hidden_size)
|
|
|
|
# out projection
|
|
self.out_proj = nn.Linear(hidden_size, hidden_size)
|
|
|
|
# rel position embeddings
|
|
self.relative_pos_embeddings = nn.Linear(hidden_size, self.num_buckets * self.num_attn_heads)
|
|
|
|
def _shape(self, tensor, seq_len, batch_size):
|
|
return paddle.transpose(
|
|
paddle.reshape(tensor, (batch_size, seq_len, self.num_attn_heads, self.head_dim)), (0, 2, 1, 3)
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
past_key_value: Optional[Tuple[Tensor]] = None,
|
|
attention_mask=None,
|
|
extended_predict_attention_mask=None,
|
|
main_relative_position_buckets=None,
|
|
predict_relative_position_buckets=None,
|
|
position_ids=None,
|
|
):
|
|
batch_size, ngram_sequence_length, hidden_size = hidden_states.shape
|
|
|
|
assert hidden_states.shape == [
|
|
batch_size,
|
|
ngram_sequence_length,
|
|
hidden_size,
|
|
], f"`hidden_states` should be of shape {batch_size, ngram_sequence_length, hidden_size}, but is of shape {hidden_states.shape}"
|
|
|
|
# project
|
|
query_states = self.query_proj(hidden_states)
|
|
key_states = self.key_proj(hidden_states)
|
|
value_states = self.value_proj(hidden_states)
|
|
|
|
# normalize
|
|
query_states = query_states / (self.head_dim**0.5)
|
|
|
|
# reshape
|
|
query_states = self._shape(query_states, ngram_sequence_length, batch_size)
|
|
key_states = self._shape(key_states, -1, batch_size)
|
|
value_states = self._shape(value_states, -1, batch_size)
|
|
|
|
proj_shape = (batch_size * self.num_attn_heads, -1, self.head_dim)
|
|
|
|
query_states = paddle.reshape(query_states, proj_shape)
|
|
key_states = paddle.reshape(key_states, proj_shape)
|
|
value_states = paddle.reshape(value_states, proj_shape)
|
|
|
|
# chunk into main stream and predict stream
|
|
hidden_states_list = paddle.chunk(hidden_states, 1 + self.ngram, axis=1)
|
|
|
|
query_states_list = paddle.chunk(query_states, 1 + self.ngram, axis=1)
|
|
key_states_list = paddle.chunk(key_states, 1 + self.ngram, axis=1)
|
|
value_states_list = paddle.chunk(value_states, 1 + self.ngram, axis=1)
|
|
|
|
main_hidden_states, hidden_states_predict_list = hidden_states_list[0], hidden_states_list[1:]
|
|
main_query_states, predict_query_states_list = query_states_list[0], query_states_list[1:]
|
|
main_key_states, predict_key_states_list = key_states_list[0], key_states_list[1:]
|
|
main_value_states, predict_value_states_list = value_states_list[0], value_states_list[1:]
|
|
|
|
# saved states are stored with shape (batch_size, num_attn_heads, seq_len, head_dim)
|
|
if past_key_value is not None:
|
|
prev_main_key_states = past_key_value[0].reshape([batch_size * self.num_attn_heads, -1, self.head_dim])
|
|
main_key_states = paddle.concat((prev_main_key_states, main_key_states), axis=1)
|
|
prev_main_value_states = past_key_value[1].reshape([batch_size * self.num_attn_heads, -1, self.head_dim])
|
|
main_value_states = paddle.concat((prev_main_value_states, main_value_states), axis=1)
|
|
|
|
# Update cache
|
|
past_key_value = (
|
|
paddle.reshape(main_key_states, (batch_size, self.num_attn_heads, -1, self.head_dim)),
|
|
paddle.reshape(main_value_states, (batch_size, self.num_attn_heads, -1, self.head_dim)),
|
|
)
|
|
|
|
# get seq_length of main stream only
|
|
sequence_length = ngram_sequence_length // (1 + self.ngram)
|
|
|
|
# MAIN-STREAM
|
|
# main attn weights
|
|
main_attn_weights = paddle.bmm(main_query_states, paddle.transpose(main_key_states, (0, 2, 1)))
|
|
|
|
# retrieve relative position embeddings for each layer -> see paper for more details
|
|
main_relative_pos_embeddings = self.get_main_relative_pos_embeddings(
|
|
main_hidden_states, main_attn_weights, position_ids, main_relative_position_buckets
|
|
)
|
|
|
|
main_attn_weights = main_attn_weights + main_relative_pos_embeddings
|
|
|
|
if attention_mask is not None:
|
|
main_attn_weights = main_attn_weights + attention_mask
|
|
|
|
main_attn_probs = F.softmax(main_attn_weights, axis=-1, dtype=main_attn_weights.dtype)
|
|
|
|
main_attn_probs = F.dropout(main_attn_probs, p=self.attention_dropout, training=self.training)
|
|
# project to attn_output
|
|
main_attn_output = paddle.bmm(main_attn_probs, main_value_states)
|
|
|
|
# reshape so that num_heads dim is merged into last `head_dim` axis
|
|
main_attn_output = paddle.reshape(
|
|
paddle.transpose(
|
|
paddle.reshape(main_attn_output, (batch_size, self.num_attn_heads, sequence_length, self.head_dim)),
|
|
(0, 2, 1, 3),
|
|
),
|
|
(batch_size, 1, sequence_length, hidden_size),
|
|
)
|
|
main_attn_output = self.out_proj(main_attn_output)
|
|
|
|
# PREDICT-STREAM
|
|
# [ngram, B*head, T, c]
|
|
predict_query_states = paddle.reshape(
|
|
paddle.concat(predict_query_states_list, axis=0), (self.ngram, -1, sequence_length, self.head_dim)
|
|
)
|
|
# [ngram, B*head, 2*T, c]
|
|
predict_key_states = paddle.concat(
|
|
[
|
|
paddle.unsqueeze(paddle.concat([main_key_states, key], axis=1), axis=0)
|
|
for key in predict_key_states_list
|
|
],
|
|
axis=0,
|
|
)
|
|
|
|
# [ngram, T, B, C]
|
|
predict_hidden_states = paddle.reshape(
|
|
paddle.concat(hidden_states_predict_list, axis=0), (self.ngram, sequence_length, batch_size, hidden_size)
|
|
)
|
|
|
|
# [ngram, B*head, 2*T, c]
|
|
predict_value_states = paddle.concat(
|
|
[
|
|
paddle.unsqueeze(paddle.concat([main_value_states, v_p], axis=1), axis=0)
|
|
for v_p in predict_value_states_list
|
|
],
|
|
axis=0,
|
|
)
|
|
|
|
# [ngram, B*head, T, 2*T]
|
|
predict_attn_weights = paddle.einsum("nbtc,nbsc->nbts", predict_query_states, predict_key_states)
|
|
|
|
# [ngram, B*head, T, S]
|
|
# retrieve relative position embeddings for each layer -> see paper for more details
|
|
predict_relative_pos_embeddings = self.get_predict_relative_pos_embeddings(
|
|
predict_hidden_states, predict_attn_weights, position_ids, predict_relative_position_buckets
|
|
)
|
|
|
|
# [ngram, B*head, T, 2*T]
|
|
predict_attn_weights = predict_attn_weights + predict_relative_pos_embeddings
|
|
|
|
if extended_predict_attention_mask is not None:
|
|
predict_attn_weights = predict_attn_weights + paddle.cast(
|
|
extended_predict_attention_mask, predict_attn_weights.dtype
|
|
)
|
|
|
|
predict_attn_probs = F.softmax(predict_attn_weights, axis=-1, dtype=predict_attn_weights.dtype)
|
|
|
|
predict_attn_probs = F.dropout(predict_attn_probs, p=self.attention_dropout, training=self.training)
|
|
# project to attention output
|
|
# [ngram, B*head, T, c]
|
|
predict_attn_output = paddle.einsum("nbts,nbsc->nbtc", predict_attn_probs, predict_value_states)
|
|
|
|
# reshape so that num_heads dim is merged into last `head_dim` axis
|
|
# [ngram, B, T, C]
|
|
predict_attn_output = paddle.reshape(
|
|
paddle.transpose(
|
|
paddle.reshape(
|
|
predict_attn_output, (self.ngram, batch_size, self.num_attn_heads, sequence_length, self.head_dim)
|
|
),
|
|
(1, 0, 3, 2, 4),
|
|
),
|
|
(batch_size, self.ngram, sequence_length, hidden_size),
|
|
)
|
|
predict_attn_output = self.out_proj(predict_attn_output)
|
|
|
|
# concat to single attn output
|
|
# [B, 1+ngram*T, C]
|
|
attn_output = paddle.reshape(
|
|
paddle.concat([main_attn_output, predict_attn_output], axis=1), (batch_size, -1, hidden_size)
|
|
)
|
|
# reshape into better form for `config.output_attentions`
|
|
main_attn_probs = paddle.reshape(main_attn_probs, (batch_size, self.num_attn_heads, sequence_length, -1))
|
|
predict_attn_probs = paddle.transpose(
|
|
paddle.reshape(predict_attn_probs, (self.ngram, batch_size, self.num_attn_heads, sequence_length, -1)),
|
|
(1, 0, 2, 3, 4),
|
|
)
|
|
|
|
attn_output = F.dropout(attn_output, p=self.dropout, training=self.training)
|
|
|
|
return attn_output, main_attn_probs, predict_attn_probs, past_key_value
|
|
|
|
def get_main_relative_pos_embeddings(
|
|
self, hidden_states, attn_weights, position_ids, main_relative_position_buckets
|
|
):
|
|
# input hidden_states [B,T,C], input attn_weights [T*head,T,S], input position_ids [B,T] or [1,1]
|
|
|
|
if main_relative_position_buckets is None:
|
|
batch_size, sequence_length = hidden_states.shape[:2]
|
|
relative_positions = paddle.tile(
|
|
paddle.unsqueeze(paddle.unsqueeze(paddle.arange(1, attn_weights.shape[-1] + 1), axis=0), axis=0),
|
|
repeat_times=[batch_size, sequence_length, 1],
|
|
)
|
|
relative_positions = relative_positions - paddle.tile(
|
|
paddle.unsqueeze(position_ids, axis=0), repeat_times=[batch_size, sequence_length, 1]
|
|
) # [B, T, s]
|
|
main_relative_position_buckets = compute_relative_buckets(
|
|
self.num_buckets, self.relative_max_distance, relative_positions, False
|
|
)
|
|
|
|
rel_pos_embeddings = self.relative_pos_embeddings(hidden_states) # [B,T,Buckets*head]
|
|
rel_pos_embeddings = paddle.transpose(
|
|
paddle.reshape(
|
|
rel_pos_embeddings, (rel_pos_embeddings.shape[:2] + [self.num_buckets, self.num_attn_heads])
|
|
),
|
|
(0, 3, 1, 2),
|
|
) # [B,T,Buckets,head]
|
|
rel_pos_embeddings = rel_pos_embeddings.reshape(attn_weights.shape[:2] + [-1]) # [B*head,T,Buckets]
|
|
|
|
main_relative_position_buckets = paddle.cast(
|
|
paddle.reshape(
|
|
paddle.tile(main_relative_position_buckets, repeat_times=[1, self.num_attn_heads, 1]),
|
|
(-1, main_relative_position_buckets.shape[-1]),
|
|
),
|
|
dtype=paddle.int64,
|
|
) # [B*head*T, T]
|
|
rel_pos_embeddings = paddle.reshape(
|
|
rel_pos_embeddings, (-1, rel_pos_embeddings.shape[-1])
|
|
) # [B*head*T,Buckets]
|
|
|
|
main_relative_position_buckets_index = paddle.tile(
|
|
main_relative_position_buckets.unsqueeze(2), repeat_times=[1, 1, 2]
|
|
)
|
|
main_relative_position_buckets_index[:, :, 0] = paddle.tile(
|
|
paddle.arange(0, main_relative_position_buckets_index.shape[0]).unsqueeze(1),
|
|
repeat_times=[1, main_relative_position_buckets_index.shape[1]],
|
|
)
|
|
|
|
main_relative_pos_embeddings = paddle.reshape(
|
|
paddle.gather_nd(rel_pos_embeddings, index=main_relative_position_buckets_index),
|
|
(attn_weights.shape[:2] + [-1]),
|
|
)
|
|
return main_relative_pos_embeddings
|
|
|
|
def get_predict_relative_pos_embeddings(
|
|
self, hidden_states, attn_weights, position_ids, predict_relative_position_buckets
|
|
):
|
|
# input hidden_states [ngram, T,B,C],
|
|
# input attn_weights [ngram, B*head,T,S],
|
|
# input position_ids [B,T] or [1,1],
|
|
# input predict_relative_position_buckets [B,T, 2*T] or None
|
|
sequence_length, batch_size = hidden_states.shape[1:3]
|
|
|
|
if predict_relative_position_buckets is None:
|
|
key_sequence_length = attn_weights.shape[-1]
|
|
assert (
|
|
position_ids[0][0] == key_sequence_length - 1
|
|
), "`position_ids` are incorrect. They should be of the format 1 2 3 4 5 ... (key_sequence_length - 1)"
|
|
relative_positions = paddle.tile(
|
|
paddle.unsqueeze(paddle.unsqueeze(paddle.arange(0, key_sequence_length), axis=0), axis=0),
|
|
repeat_times=[batch_size, sequence_length, 1],
|
|
)
|
|
|
|
relative_positions = relative_positions - paddle.tile(
|
|
paddle.unsqueeze(position_ids, axis=0), repeat_times=[batch_size, sequence_length, 1]
|
|
)
|
|
predict_relative_position_buckets = compute_relative_buckets(
|
|
self.num_buckets, self.relative_max_distance, relative_positions, False
|
|
)
|
|
|
|
hidden_states = paddle.transpose(hidden_states, (0, 2, 1, 3)) # [ngram, B, T, C]
|
|
rel_pos_embeddings = paddle.reshape(
|
|
self.relative_pos_embeddings(hidden_states),
|
|
hidden_states.shape[:-1] + [self.num_buckets, self.num_attn_heads],
|
|
) # [ngram, B, T, bucket, head]
|
|
rel_pos_embeddings = paddle.reshape(
|
|
paddle.transpose(rel_pos_embeddings, (0, 1, 4, 2, 3)),
|
|
(self.ngram * batch_size * self.num_attn_heads, sequence_length, -1),
|
|
) # [ngram*B*head, T, bucket]
|
|
|
|
predict_relative_position_buckets = paddle.tile(
|
|
paddle.unsqueeze(predict_relative_position_buckets, axis=0),
|
|
repeat_times=[self.ngram, 1, self.num_attn_heads, 1],
|
|
) # [ngram, B, head*T, S]
|
|
|
|
rel_pos_embeddings = paddle.reshape(rel_pos_embeddings, (-1, rel_pos_embeddings.shape[-1]))
|
|
predict_relative_position_buckets = paddle.cast(
|
|
paddle.reshape(predict_relative_position_buckets, (-1, predict_relative_position_buckets.shape[-1])),
|
|
dtype=paddle.int64,
|
|
) # [ngram*B*head*T, S]
|
|
|
|
predict_relative_position_buckets_index = paddle.tile(
|
|
predict_relative_position_buckets.unsqueeze(2), repeat_times=[1, 1, 2]
|
|
)
|
|
predict_relative_position_buckets_index[:, :, 0] = paddle.tile(
|
|
paddle.arange(0, predict_relative_position_buckets_index.shape[0]).unsqueeze(1),
|
|
repeat_times=[1, predict_relative_position_buckets_index.shape[1]],
|
|
)
|
|
|
|
predict_relative_pos_embeddings = paddle.reshape(
|
|
paddle.gather_nd(rel_pos_embeddings, index=predict_relative_position_buckets_index),
|
|
(self.ngram, batch_size * self.num_attn_heads, sequence_length, -1),
|
|
) # [ngram, B*head, T, S]
|
|
|
|
return predict_relative_pos_embeddings
|
|
|
|
|
|
class ProphetNetEncoderLayer(Layer):
|
|
"""
|
|
Encoder block for Prophetnet
|
|
"""
|
|
|
|
def __init__(self, config: ProphetNetConfig):
|
|
super(ProphetNetEncoderLayer, self).__init__()
|
|
# 1st residual block
|
|
self.self_attn = ProphetNetAttention(
|
|
config.hidden_size, config.attention_dropout, config.dropout, config.num_encoder_attention_heads
|
|
)
|
|
self.self_attn_layer_norm = nn.LayerNorm(config.hidden_size)
|
|
|
|
# 2nd residual block
|
|
self.feed_forward = ProphetNetFeedForward(
|
|
config.hidden_size,
|
|
config.activation_function,
|
|
config.activation_dropout,
|
|
config.dropout,
|
|
config.encoder_ffn_dim,
|
|
)
|
|
self.feed_forward_layer_norm = nn.LayerNorm(config.hidden_size)
|
|
|
|
def forward(self, hidden_states, attention_mask):
|
|
# 1st residual block
|
|
attention_output, _ = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask)
|
|
hidden_states = self.self_attn_layer_norm(attention_output + hidden_states)
|
|
|
|
# 2nd residual block
|
|
feed_forward_output = self.feed_forward(hidden_states)
|
|
hidden_states = self.feed_forward_layer_norm(feed_forward_output + hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
class ProphetNetDecoderLayer(Layer):
|
|
"""
|
|
Decoder block for Prophetnet
|
|
"""
|
|
|
|
def __init__(self, config: ProphetNetConfig):
|
|
super(ProphetNetDecoderLayer, self).__init__()
|
|
# 1st residual block
|
|
self.self_attn = ProphetNetNgramSelfAttention(
|
|
config.hidden_size,
|
|
config.num_buckets,
|
|
config.relative_max_distance,
|
|
config.num_decoder_attention_heads,
|
|
config.dropout,
|
|
config.attention_dropout,
|
|
config.ngram,
|
|
)
|
|
self.self_attn_layer_norm = nn.LayerNorm(config.hidden_size)
|
|
|
|
# 2nd residual block
|
|
if config.add_cross_attention:
|
|
self.cross_attn = ProphetNetAttention(
|
|
config.hidden_size, config.attention_dropout, config.dropout, config.num_decoder_attention_heads
|
|
)
|
|
self.cross_attn_layer_norm = nn.LayerNorm(config.hidden_size)
|
|
|
|
# 3rd residual block
|
|
self.feed_forward = ProphetNetFeedForward(
|
|
config.hidden_size,
|
|
config.activation_function,
|
|
config.activation_dropout,
|
|
config.dropout,
|
|
config.decoder_ffn_dim,
|
|
)
|
|
self.feed_forward_layer_norm = nn.LayerNorm(config.hidden_size)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
attention_mask=None,
|
|
encoder_hidden_states=None,
|
|
encoder_attn_mask=None,
|
|
extended_predict_attention_mask=None,
|
|
main_relative_position_buckets=None,
|
|
predict_relative_position_buckets=None,
|
|
position_ids=None,
|
|
past_key_value=None,
|
|
use_cache: bool = True,
|
|
):
|
|
# 1st residual block
|
|
# decoder uni-directional self-attention cached key/values tuple is at positions 1,2
|
|
self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
|
|
ngram_attention_output, self_attn_weights, self_attn_weights_ngram, present_key_value = self.self_attn(
|
|
hidden_states=hidden_states,
|
|
past_key_value=self_attn_past_key_value,
|
|
attention_mask=attention_mask,
|
|
extended_predict_attention_mask=extended_predict_attention_mask,
|
|
main_relative_position_buckets=main_relative_position_buckets,
|
|
predict_relative_position_buckets=predict_relative_position_buckets,
|
|
position_ids=position_ids,
|
|
)
|
|
hidden_states = self.self_attn_layer_norm(hidden_states + ngram_attention_output)
|
|
|
|
# cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple
|
|
cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
|
|
if encoder_hidden_states is not None:
|
|
# 2nd residual block
|
|
attention_output, cross_attn_present_key_value = self.cross_attn(
|
|
hidden_states=hidden_states,
|
|
key_value_states=encoder_hidden_states,
|
|
attention_mask=encoder_attn_mask,
|
|
past_key_value=cross_attn_past_key_value,
|
|
)
|
|
hidden_states = self.cross_attn_layer_norm(attention_output + hidden_states)
|
|
|
|
# add cross-attn to positions 3,4 of present_key_value tuple
|
|
present_key_value = present_key_value + cross_attn_present_key_value
|
|
|
|
# 3rd residual block
|
|
feed_forward_output = self.feed_forward(hidden_states)
|
|
hidden_states = self.feed_forward_layer_norm(feed_forward_output + hidden_states)
|
|
|
|
outputs = (hidden_states,)
|
|
|
|
if use_cache:
|
|
outputs += (present_key_value,)
|
|
|
|
return outputs
|
|
|
|
|
|
class ProphetNetEncoder(ProphetNetPretrainedModel):
|
|
r"""
|
|
word_embeddings (:obj:`paddle.nn.Embeddings` of shape :obj:`(config.vocab_size, config.hidden_size)`, `optional`):
|
|
The word embedding parameters. This can be used to initialize :class:`~transformers.ProphetNetEncoder` with
|
|
pre-defined word embeddings instead of randomly initialized word embeddings.
|
|
"""
|
|
|
|
def __init__(self, word_embeddings, config: ProphetNetConfig):
|
|
super(ProphetNetEncoder, self).__init__(config)
|
|
self.init_std = config.init_std
|
|
if word_embeddings is not None:
|
|
self.word_embeddings = word_embeddings
|
|
else:
|
|
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
|
|
|
|
self.position_embeddings = ProphetNetPositionalEmbeddings(config)
|
|
self.embeddings_layer_norm = nn.LayerNorm(config.hidden_size)
|
|
|
|
self.layers = nn.LayerList([ProphetNetEncoderLayer(config) for _ in range(config.num_encoder_layers)])
|
|
|
|
def forward(self, input_ids=None, attention_mask=None):
|
|
if input_ids is None:
|
|
raise ValueError("Input_ids cannot be None.")
|
|
inputs_embeds = self.word_embeddings(input_ids)
|
|
|
|
# prepare attention mask
|
|
if attention_mask is not None:
|
|
extended_attention_mask = (
|
|
paddle.tile(
|
|
1.0 - attention_mask.unsqueeze(1), repeat_times=[self.config.num_encoder_attention_heads, 1, 1]
|
|
)
|
|
) * -10000.0
|
|
extended_attention_mask = paddle.cast(extended_attention_mask, dtype=inputs_embeds.dtype)
|
|
extended_attention_mask.stop_gradient = True
|
|
else:
|
|
extended_attention_mask = None
|
|
|
|
position_embeddings, position_ids = self.position_embeddings(inputs_embeds.shape[:2])
|
|
|
|
hidden_states = inputs_embeds + position_embeddings
|
|
hidden_states = self.embeddings_layer_norm(hidden_states)
|
|
hidden_states = F.dropout(hidden_states, p=self.config.dropout, training=self.training)
|
|
|
|
for idx, encoder_layer in enumerate(self.layers):
|
|
hidden_states = encoder_layer(hidden_states, attention_mask=extended_attention_mask)
|
|
return hidden_states
|
|
|
|
|
|
class ProphetNetDecoder(ProphetNetPretrainedModel):
|
|
def __init__(self, word_embeddings, config: ProphetNetConfig):
|
|
super(ProphetNetDecoder, self).__init__(config)
|
|
self.init_std = config.init_std
|
|
self.ngram = config.ngram
|
|
self.num_buckets = config.num_buckets
|
|
self.relative_max_distance = config.relative_max_distance
|
|
self.dropout = config.dropout
|
|
self.max_target_positions = config.max_position_embeddings
|
|
self.add_cross_attention = config.add_cross_attention
|
|
if word_embeddings is not None:
|
|
self.word_embeddings = word_embeddings
|
|
else:
|
|
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
|
|
|
|
self.position_embeddings = ProphetNetPositionalEmbeddings(config)
|
|
|
|
self.ngram_embeddings = nn.Embedding(self.ngram, config.hidden_size)
|
|
self.layers = nn.LayerList([ProphetNetDecoderLayer(config) for _ in range(config.num_decoder_layers)])
|
|
self.embeddings_layer_norm = nn.LayerNorm(config.hidden_size)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids=None,
|
|
attention_mask=None,
|
|
encoder_hidden_states=None,
|
|
encoder_attention_mask=None,
|
|
past_key_values=None,
|
|
use_cache=True,
|
|
):
|
|
if input_ids is None:
|
|
raise ValueError("Decoder input_ids cannot be None.")
|
|
inputs_embeds = self.word_embeddings(input_ids)
|
|
batch_size, sequence_length = inputs_embeds.shape[:2]
|
|
|
|
main_stream_pos_embed, position_ids = self.position_embeddings(
|
|
(batch_size, sequence_length), past_key_values=past_key_values
|
|
)
|
|
|
|
if past_key_values is not None:
|
|
main_relative_position_buckets, predict_relative_position_buckets = None, None
|
|
else:
|
|
main_relative_position_buckets, predict_relative_position_buckets = self.compute_buffered_relative_buckets(
|
|
position_ids
|
|
)
|
|
predicting_stream_pos_embed = self.position_embeddings._forward(position_ids + 1)
|
|
|
|
# add position embeddings
|
|
hidden_states = inputs_embeds + main_stream_pos_embed
|
|
|
|
ngram_embeddings = self.ngram_embeddings.weight
|
|
|
|
# prepare attention mask
|
|
if past_key_values is not None:
|
|
assert (
|
|
hidden_states.shape[1] == 1
|
|
), "At the moment `use_cache` is only supported for `decoder_input_ids` of length 1"
|
|
|
|
ngram_hidden_states = [
|
|
paddle.tile(
|
|
(ngram_embeddings[ngram - 1] + predicting_stream_pos_embed), repeat_times=[batch_size, 1, 1]
|
|
)
|
|
for ngram in range(self.ngram)
|
|
]
|
|
extended_attention_mask = None
|
|
extended_predict_attention_mask = None
|
|
else:
|
|
ngram_hidden_states = [
|
|
(ngram_embeddings[ngram - 1] + predicting_stream_pos_embed) for ngram in range(self.ngram)
|
|
]
|
|
extended_attention_mask = self.prepare_attention_mask(hidden_states, attention_mask)
|
|
extended_predict_attention_mask = self.prepare_predict_attention_mask(hidden_states, attention_mask)
|
|
extended_attention_mask.stop_gradient = True
|
|
extended_predict_attention_mask.stop_gradient = True
|
|
|
|
# prepare encoder attention mask
|
|
if encoder_attention_mask is not None:
|
|
extended_encoder_attention_mask = (
|
|
1.0
|
|
- paddle.tile(
|
|
encoder_attention_mask[:, None, :], repeat_times=[self.config.num_decoder_attention_heads, 1, 1]
|
|
)
|
|
) * -10000.0
|
|
extended_encoder_attention_mask = paddle.cast(extended_encoder_attention_mask, dtype=inputs_embeds.dtype)
|
|
else:
|
|
extended_encoder_attention_mask = None
|
|
|
|
hidden_states = paddle.concat([hidden_states] + ngram_hidden_states, axis=1)
|
|
|
|
if self.embeddings_layer_norm:
|
|
hidden_states = self.embeddings_layer_norm(hidden_states)
|
|
|
|
hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
|
|
present_key_values = () if use_cache else None
|
|
|
|
for idx, decoder_layer in enumerate(self.layers):
|
|
|
|
past_key_value = past_key_values[idx] if past_key_values is not None else None
|
|
|
|
layer_outputs = decoder_layer(
|
|
hidden_states,
|
|
attention_mask=extended_attention_mask,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
encoder_attn_mask=extended_encoder_attention_mask,
|
|
extended_predict_attention_mask=extended_predict_attention_mask,
|
|
main_relative_position_buckets=main_relative_position_buckets,
|
|
predict_relative_position_buckets=predict_relative_position_buckets,
|
|
position_ids=position_ids,
|
|
past_key_value=past_key_value,
|
|
use_cache=use_cache,
|
|
)
|
|
|
|
hidden_states = layer_outputs[0]
|
|
|
|
if use_cache:
|
|
present_key_values += (layer_outputs[1],)
|
|
|
|
last_hidden_state = hidden_states[:, :sequence_length] # 1-gram
|
|
last_hidden_state_ngram = hidden_states[:, sequence_length:] if self.ngram > 0 else None # 2-gram
|
|
return tuple(v for v in [last_hidden_state, last_hidden_state_ngram, present_key_values] if v is not None)
|
|
|
|
def compute_buffered_relative_buckets(self, position_ids):
|
|
batch_size, sequence_length = position_ids.shape
|
|
|
|
if not hasattr(self, "_main_relative_buckets") or self._main_relative_buckets is None:
|
|
position_ids = paddle.tile(paddle.arange(1, self.max_target_positions + 1), repeat_times=[1, 1])
|
|
self._main_relative_buckets, self._predict_relative_buckets = compute_all_stream_relative_buckets(
|
|
self.num_buckets, self.relative_max_distance, position_ids
|
|
)
|
|
|
|
# buffer relative buckets
|
|
main_relative_buckets = paddle.tile(
|
|
self._main_relative_buckets[:, :sequence_length, :sequence_length], repeat_times=[batch_size, 1, 1]
|
|
)
|
|
predict_relative_buckets = paddle.tile(
|
|
paddle.concat(
|
|
[
|
|
self._predict_relative_buckets[:, :sequence_length, :sequence_length],
|
|
self._predict_relative_buckets[
|
|
:, :sequence_length, self.max_target_positions : self.max_target_positions + sequence_length
|
|
],
|
|
],
|
|
axis=2,
|
|
),
|
|
repeat_times=[batch_size, 1, 1],
|
|
)
|
|
|
|
return main_relative_buckets, predict_relative_buckets
|
|
|
|
def prepare_attention_mask(self, hidden_states, attention_mask):
|
|
batch_size, seq_length = hidden_states.shape[:2]
|
|
|
|
# get causal mask
|
|
if not hasattr(self, "_causal_mask") or self._causal_mask is None:
|
|
causal_mask = paddle.full(
|
|
(self.max_target_positions, self.max_target_positions), -float("inf"), dtype=hidden_states.dtype
|
|
)
|
|
self._causal_mask = paddle.triu(causal_mask, 1)
|
|
extended_causal_mask = paddle.expand(
|
|
self._causal_mask[:seq_length, :seq_length].unsqueeze(0), shape=[batch_size, seq_length, seq_length]
|
|
)
|
|
|
|
# add usual attention mask
|
|
if attention_mask is not None:
|
|
extended_attention_mask = (1.0 - attention_mask.unsqueeze(1)) * -10000.0
|
|
extended_attention_mask = extended_causal_mask + extended_attention_mask
|
|
else:
|
|
extended_attention_mask = extended_causal_mask
|
|
return paddle.cast(
|
|
paddle.tile(extended_attention_mask, repeat_times=[self.config.num_decoder_attention_heads, 1, 1]),
|
|
dtype=hidden_states.dtype,
|
|
)
|
|
|
|
def prepare_predict_attention_mask(self, hidden_states, attention_mask):
|
|
batch_size, seq_length = hidden_states.shape[:2]
|
|
|
|
# get causal mask
|
|
if not hasattr(self, "_predict_causal_mask") or self._predict_causal_mask is None:
|
|
self._predict_causal_mask = ngram_attention_bias(
|
|
self.max_target_positions, self.ngram, hidden_states.dtype
|
|
)
|
|
predict_causal_mask = paddle.concat(
|
|
[
|
|
self._predict_causal_mask[:, :seq_length, :seq_length],
|
|
self._predict_causal_mask[
|
|
:, :seq_length, self.max_target_positions : self.max_target_positions + seq_length
|
|
],
|
|
],
|
|
axis=-1,
|
|
)
|
|
extended_predict_causal_mask = paddle.expand(
|
|
predict_causal_mask[:, None, :, :],
|
|
shape=predict_causal_mask.shape[:1] + [batch_size] + predict_causal_mask.shape[1:],
|
|
)
|
|
|
|
# add usual attention mask
|
|
if attention_mask is not None:
|
|
extended_attention_mask = (1.0 - attention_mask[None, :, None, :]) * -10000.0
|
|
extended_attention_mask = extended_attention_mask.expand((self.ngram, batch_size, seq_length, seq_length))
|
|
# predicted stream attention_mask should always be 0
|
|
extended_attention_mask = paddle.concat(
|
|
[extended_attention_mask, paddle.zeros_like(extended_attention_mask)], axis=-1
|
|
)
|
|
extended_predict_attention_mask = extended_predict_causal_mask + extended_attention_mask
|
|
else:
|
|
extended_predict_attention_mask = extended_predict_causal_mask
|
|
return paddle.cast(
|
|
extended_predict_attention_mask.tile([1, self.config.num_decoder_attention_heads, 1, 1]),
|
|
dtype=hidden_states.dtype,
|
|
)
|
|
|
|
|
|
@register_base_model
|
|
class ProphetNetModel(ProphetNetPretrainedModel):
|
|
def __init__(self, config: ProphetNetConfig):
|
|
super(ProphetNetModel, self).__init__(config)
|
|
self.init_std = config.init_std
|
|
self.eps = config.eps
|
|
self.pad_token_id = config.pad_token_id
|
|
self.disable_ngram_loss = config.disable_ngram_loss
|
|
self.decoder_start_token_id = config.decoder_start_token_id
|
|
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
|
|
|
|
self.encoder = ProphetNetEncoder(self.word_embeddings, config)
|
|
|
|
self.decoder = ProphetNetDecoder(self.word_embeddings, config)
|
|
|
|
def get_encoder(self):
|
|
return self.encoder
|
|
|
|
def get_decoder(self):
|
|
return self.decoder
|
|
|
|
def get_input_embeddings(self):
|
|
return self.word_embeddings
|
|
|
|
def set_input_embeddings(self, value):
|
|
self.word_embeddings = value
|
|
|
|
def forward(
|
|
self,
|
|
input_ids=None,
|
|
attention_mask=None,
|
|
decoder_input_ids=None,
|
|
decoder_attention_mask=None,
|
|
encoder_output: Optional[Tuple] = None,
|
|
use_cache=True,
|
|
past_key_values=None,
|
|
):
|
|
if attention_mask is None:
|
|
assert input_ids is not None, "input_ids should be " "specified when generating attention_mask"
|
|
attention_mask = paddle.cast(input_ids != self.pad_token_id, dtype=paddle.get_default_dtype())
|
|
|
|
if decoder_attention_mask is None:
|
|
assert decoder_input_ids is not None, (
|
|
"decoder_input_ids should be " "specified when generating decoder_attention_mask"
|
|
)
|
|
decoder_attention_mask = paddle.cast(
|
|
decoder_input_ids != self.pad_token_id, dtype=paddle.get_default_dtype()
|
|
)
|
|
if encoder_output is None:
|
|
encoder_output = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
|
|
decoder_outputs = self.decoder(
|
|
input_ids=decoder_input_ids,
|
|
attention_mask=decoder_attention_mask,
|
|
encoder_hidden_states=encoder_output,
|
|
encoder_attention_mask=attention_mask,
|
|
use_cache=use_cache,
|
|
past_key_values=past_key_values,
|
|
)
|
|
return decoder_outputs + (encoder_output,)
|
|
|
|
|
|
class Linear_wo_bias(Layer):
|
|
def __init__(self, in_features, out_features, weight_attr=None, name=None):
|
|
super(Linear_wo_bias, self).__init__()
|
|
self._dtype = self._helper.get_default_dtype()
|
|
self._weight_attr = weight_attr
|
|
self.weight = self.create_parameter(
|
|
shape=[in_features, out_features], attr=self._weight_attr, dtype=self._dtype, is_bias=False
|
|
)
|
|
self.name = name
|
|
|
|
def forward(self, input):
|
|
out = F.linear(x=input, weight=self.weight, name=self.name)
|
|
return out
|
|
|
|
def extra_repr(self):
|
|
name_str = ", name={}".format(self.name) if self.name else ""
|
|
return "in_features={}, out_features={}, dtype={}{}".format(
|
|
self.weight.shape[0], self.weight.shape[1], self._dtype, name_str
|
|
)
|
|
|
|
|
|
class ProphetNetForConditionalGeneration(ProphetNetPretrainedModel):
|
|
def __init__(self, config: ProphetNetConfig):
|
|
super(ProphetNetForConditionalGeneration, self).__init__(config)
|
|
self.prophetnet = ProphetNetModel(config)
|
|
self.padding_idx = self.prophetnet.word_embeddings._padding_idx
|
|
|
|
self.lm_head = Linear_wo_bias(config.hidden_size, config.vocab_size)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids=None,
|
|
attention_mask=None,
|
|
decoder_input_ids=None,
|
|
decoder_attention_mask=None,
|
|
encoder_output=None,
|
|
labels=None,
|
|
use_cache=True,
|
|
past_key_values=None,
|
|
):
|
|
if labels is not None and decoder_input_ids is None:
|
|
# get decoder inputs from shifting lm labels to the right
|
|
decoder_input_ids = self._shift_right(labels)
|
|
outputs = self.prophetnet(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
decoder_input_ids=decoder_input_ids,
|
|
decoder_attention_mask=decoder_attention_mask,
|
|
encoder_output=encoder_output,
|
|
use_cache=use_cache,
|
|
past_key_values=past_key_values,
|
|
)
|
|
|
|
batch_size, sequence_length = decoder_input_ids.shape
|
|
|
|
predicting_streams = paddle.reshape(outputs[1], (batch_size, self.config.ngram, sequence_length, -1))
|
|
predict_logits = self.lm_head(predicting_streams)
|
|
|
|
logits = predict_logits[:, 0]
|
|
if use_cache:
|
|
past_key_values = outputs[2]
|
|
return logits, past_key_values, predict_logits
|
|
else:
|
|
return logits, predict_logits
|
|
|
|
def prepare_inputs_for_generation(
|
|
self,
|
|
decoder_input_ids,
|
|
attention_mask=None,
|
|
decoder_attention_mask=None,
|
|
cache=None,
|
|
use_cache=None,
|
|
encoder_output=None,
|
|
):
|
|
assert encoder_output is not None, "`encoder_output` have to be passed for generation."
|
|
if cache is not None:
|
|
decoder_input_ids = decoder_input_ids[:, -1].unsqueeze(-1)
|
|
|
|
# first step, decoder_cached_states are empty
|
|
return {
|
|
"input_ids": None, # encoder_outputs is defined. input_ids not needed
|
|
"decoder_input_ids": decoder_input_ids,
|
|
"encoder_output": encoder_output,
|
|
"decoder_attention_mask": decoder_attention_mask,
|
|
"attention_mask": attention_mask,
|
|
"use_cache": use_cache,
|
|
"past_key_values": cache,
|
|
}
|
|
|
|
def prepare_decoder_input_ids_from_labels(self, labels):
|
|
return self._shift_right(labels)
|
|
|
|
def get_encoder(self):
|
|
return self.prophetnet.encoder
|
|
|
|
def get_decoder(self):
|
|
return self.prophetnet.decoder
|
|
|
|
def __getattr__(self, name):
|
|
try:
|
|
return super().__getattr__(name)
|
|
except AttributeError:
|
|
return getattr(getattr(self, self.base_model_prefix), name)
|