-
Notifications
You must be signed in to change notification settings - Fork 0
/
base_model.py
82 lines (66 loc) · 3.14 KB
/
base_model.py
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
from torch import Tensor
import torch
import torch.nn as nn
from torch.nn import Transformer
import math
def generate_square_subsequent_mask(sz):
mask = (torch.triu(torch.ones((sz, sz))) == 1).transpose(0, 1)
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
return mask
class PositionalEncoding(nn.Module):
def __init__(self,
emb_size: int,
dropout: float,
maxlen: int = 5000):
super(PositionalEncoding, self).__init__()
den = torch.exp(- torch.arange(0, emb_size, 2)* math.log(10000) / emb_size)
pos = torch.arange(0, maxlen).reshape(maxlen, 1)
pos_embedding = torch.zeros((maxlen, emb_size))
pos_embedding[:, 0::2] = torch.sin(pos * den)
pos_embedding[:, 1::2] = torch.cos(pos * den)
pos_embedding = pos_embedding.unsqueeze(-2)
self.dropout = nn.Dropout(dropout)
self.register_buffer('pos_embedding', pos_embedding)
def forward(self, token_embedding: Tensor):
return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])
class KnowledgeModel(nn.Module):
def __init__(self,
num_encoder_layers: int,
num_decoder_layers: int,
emb_size: int,
nhead: int,
dim_feedforward: int = 512,
dropout: float = 0.1):
super(KnowledgeModel, self).__init__()
self.encoder_proj = nn.Linear(768, 256)
self.transformer = Transformer(d_model=emb_size,
nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dim_feedforward=dim_feedforward,
dropout=dropout)
self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)
self.generator = nn.Linear(emb_size, 1)
self._initialize_params()
def _initialize_params(self):
for p in self.transformer.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
def forward(self,
src_emb: Tensor,
tgt_emb: Tensor,
src_mask: Tensor,
tgt_mask: Tensor,
src_padding_mask: Tensor,
tgt_padding_mask: Tensor,
memory_key_padding_mask: Tensor):
src_emb = self.encoder_proj(src_emb)
src_emb = self.positional_encoding(src_emb)
tgt_emb = self.positional_encoding(tgt_emb)
outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None,
src_padding_mask, tgt_padding_mask, memory_key_padding_mask)
return self.generator(outs)
def encode(self, src_emb: Tensor, src_mask: Tensor):
return self.transformer.encoder(self.positional_encoding(src_emb), src_mask)
def decode(self, tgt_emb: Tensor, memory: Tensor, tgt_mask: Tensor):
return self.transformer.decoder(self.positional_encoding(tgt_emb), memory, tgt_mask)