Spaces:
Runtime error
Runtime error
File size: 1,158 Bytes
7900c16 |
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 |
import torch.nn as nn
import torch
from tencentpretrain.layers.layer_norm import LayerNorm
class Embedding(nn.Module):
def __init__(self, args):
super(Embedding, self).__init__()
self.embedding_name_list = []
self.dropout = nn.Dropout(args.dropout)
self.remove_embedding_layernorm = args.remove_embedding_layernorm
if not self.remove_embedding_layernorm and "dual" not in args.embedding:
self.layer_norm = LayerNorm(args.emb_size)
def update(self, embedding, embedding_name):
setattr(self, embedding_name, embedding)
self.embedding_name_list.append(embedding_name)
def forward(self, src, seg):
if self.embedding_name_list[0] == "dual":
return self.dual(src, seg)
for i, embedding_name in enumerate(self.embedding_name_list):
embedding = getattr(self, embedding_name)
if i == 0:
emb = embedding(src, seg)
else:
emb += embedding(src, seg)
if not self.remove_embedding_layernorm:
emb = self.layer_norm(emb)
emb = self.dropout(emb)
return emb
|