67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Embedder(nn.Module):
|
|
"""
|
|
Composite embedding layer.
|
|
"""
|
|
|
|
def __init__(self,
|
|
hidden_dim,
|
|
num_token_embeddings,
|
|
num_pos_embeddings,
|
|
num_type_embeddings,
|
|
num_turn_embeddings,
|
|
padding_idx=None,
|
|
dropout=0.1,
|
|
pos_trainable=False):
|
|
super(Embedder, self).__init__()
|
|
|
|
self.token_embedding = nn.Embedding(num_token_embeddings, hidden_dim)
|
|
self.pos_embedding = nn.Embedding(num_pos_embeddings, hidden_dim)
|
|
self.pos_embedding.weight.requires_grad = pos_trainable
|
|
self.type_embedding = nn.Embedding(num_type_embeddings, hidden_dim)
|
|
self.turn_embedding = nn.Embedding(num_turn_embeddings, hidden_dim)
|
|
self.dropout_layer = nn.Dropout(p=dropout)
|
|
|
|
# follow the default xavier_uniform initializer in paddle version
|
|
# otherwise, there are bugs for dec_probs computation in weight typing setting
|
|
# default norm initializer in nn.Embedding in pytorch, which samples larger values
|
|
nn.init.xavier_uniform_(self.token_embedding.weight)
|
|
nn.init.xavier_uniform_(self.pos_embedding.weight)
|
|
nn.init.xavier_uniform_(self.type_embedding.weight)
|
|
nn.init.xavier_uniform_(self.turn_embedding.weight)
|
|
return
|
|
|
|
def forward(self, token_inp, pos_inp=None, type_inp=None, turn_inp=None):
|
|
embed = self.token_embedding(token_inp)
|
|
if pos_inp is not None:
|
|
embed += self.pos_embedding(pos_inp)
|
|
if type_inp is not None:
|
|
embed += self.type_embedding(type_inp)
|
|
if turn_inp is not None:
|
|
embed += self.turn_embedding(turn_inp)
|
|
embed = self.dropout_layer(embed)
|
|
return embed
|
|
|
|
|
|
def main():
|
|
import numpy as np
|
|
|
|
model = Embedder(10, 20, 20, 20, 20)
|
|
token_inp = torch.tensor(
|
|
np.random.randint(0, 19, [10, 10]).astype('int64'))
|
|
pos_inp = torch.tensor(np.random.randint(0, 19, [10, 10]).astype('int64'))
|
|
type_inp = torch.tensor(np.random.randint(0, 19, [10, 10]).astype('int64'))
|
|
turn_inp = torch.tensor(np.random.randint(0, 19, [10, 10]).astype('int64'))
|
|
out = model(token_inp, pos_inp, type_inp, turn_inp)
|
|
print(out)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|