| 72 | |
| 73 | |
| 74 | class TransformerInterEncoder(nn.Module): |
| 75 | def __init__(self, d_model, d_ff, heads, dropout, num_inter_layers=0): |
| 76 | super(TransformerInterEncoder, self).__init__() |
| 77 | self.d_model = d_model |
| 78 | self.num_inter_layers = num_inter_layers |
| 79 | self.pos_emb = PositionalEncoding(dropout, d_model) |
| 80 | self.transformer_inter = nn.ModuleList( |
| 81 | [TransformerEncoderLayer(d_model, heads, d_ff, dropout) |
| 82 | for _ in range(num_inter_layers)]) |
| 83 | self.dropout = nn.Dropout(dropout) |
| 84 | self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) |
| 85 | self.wo = nn.Linear(d_model, 1, bias=True) |
| 86 | self.sigmoid = nn.Sigmoid() |
| 87 | # self.last_status = None |
| 88 | |
| 89 | def forward(self, top_vecs, mask): |
| 90 | """ See :obj:`EncoderBase.forward()`""" |
| 91 | |
| 92 | batch_size, n_sents = top_vecs.size(0), top_vecs.size(1) |
| 93 | pos_emb = self.pos_emb.pe[:, :n_sents] |
| 94 | x = top_vecs * mask[:, :, None].float() |
| 95 | x = x + pos_emb |
| 96 | |
| 97 | for i in range(self.num_inter_layers): |
| 98 | x = self.transformer_inter[i](i, x, x, ~ mask) # all_sents * max_tokens * dim |
| 99 | |
| 100 | x = self.layer_norm(x) |
| 101 | # print(x.size()) |
| 102 | # TODO: last status vector return this vector |
| 103 | # self.last_status = x |
| 104 | sent_scores = self.sigmoid(self.wo(x)) |
| 105 | sent_scores = sent_scores.squeeze(-1) * mask.float() |
| 106 | return sent_scores |
| 107 | |
| 108 | |
| 109 | class RNNEncoder(nn.Module): |