File size: 1,629 Bytes
101e5d8 9d1e66f 101e5d8 ae1bf72 101e5d8 ae1bf72 101e5d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
# coding=utf-8
# author: xusong <xusong28@jd.com>
# time: 2022/3/3 14:18
"""
?
"""
import torch
from torch import nn
from transformers.models.bart import modeling_bart
class KPlugLearnedPositionalEmbedding(nn.Embedding):
"""
This module learns positional embeddings up to a fixed maximum size.
"""
def __init__(self, num_embeddings: int, embedding_dim: int):
# Bart is set up so that if padding_idx is specified then offset the embedding ids by 2
# and adjust num_embeddings appropriately. Other models don't have this hack
self.offset = 2
super().__init__(num_embeddings + self.offset, embedding_dim)
### 4.21.1 之后的版本
# def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):
# """`input_ids' shape is expected to be [bsz x seqlen]."""
#
# bsz, seq_len = input_ids.shape[:2]
# positions = torch.arange(
# past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
# ).expand(bsz, -1)
#
# return super().forward(positions + self.offset)
def forward(self, input_ids_shape: torch.Size, past_key_values_length: int = 0):
"""`input_ids_shape` is expected to be [bsz x seqlen]."""
bsz, seq_len = input_ids_shape[:2]
positions = torch.arange(
past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
)
return super().forward(positions + self.offset)
modeling_bart.BartLearnedPositionalEmbedding = KPlugLearnedPositionalEmbedding
|