Spaces:
Sleeping
Sleeping
File size: 3,121 Bytes
383bfb8 |
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 |
import torch
import torch.nn as nn
class Score(nn.Module):
def __init__(self, input_size, hidden_size):
super(Score, self).__init__()
self.attn = nn.Linear(hidden_size + input_size, hidden_size)
self.score = nn.Linear(hidden_size, 1, bias=False)
def forward(self, hidden, candi_embeddings, candi_mask=None):
'''
Arguments:
hidden: B x 1 x 2H
candi_embeddings: B x candi_size x H
candi_mask: B x candi_size
Return:
score: B x candi_size
'''
hidden = hidden.repeat(1, candi_embeddings.size(1), 1) # B x candi_size x H
# For each position of encoder outputs
energy_in = torch.cat((hidden, candi_embeddings), 2) # B x candi_size x 3H
score = self.score(torch.tanh(self.attn(energy_in))).squeeze(-1) # B x candi_size
if candi_mask is not None:
score = score.masked_fill_(~candi_mask, -1e12)
return score
class Attn(nn.Module):
def __init__(self, input_size, hidden_size):
super(Attn, self).__init__()
self.attn = nn.Linear(hidden_size + input_size, hidden_size)
self.score = nn.Linear(hidden_size, 1, bias=False)
def forward(self, hidden, encoder_outputs, seq_mask=None):
'''
Arguments:
hidden: B x 1 x H (q)
encoder_outputs: B x S x H
seq_mask: B x S
Return:
attn_energies: B x S
'''
hidden = hidden.repeat(1, encoder_outputs.size(1), 1) # B x S x H
energy_in = torch.cat((hidden, encoder_outputs), 2) # B x S x 2H
score_feature = torch.tanh(self.attn(energy_in)) # B x S x H
attn_energies = self.score(score_feature).squeeze(-1) # B x S
if seq_mask is not None:
attn_energies = attn_energies.masked_fill_(~seq_mask, -1e12)
attn_energies = nn.functional.softmax(attn_energies, dim=1) # B x S
return attn_energies
class Score_Multi(nn.Module):
def __init__(self, input_size, hidden_size):
super(Score_Multi, self).__init__()
self.attn = nn.Linear(hidden_size + input_size, hidden_size)
self.score = nn.Linear(hidden_size, 1, bias=False)
def forward(self, hidden, candi_embeddings, candi_mask=None):
'''
Arguments:
hidden: B x S x H
candi_embeddings: B x candi_size x H
candi_mask: B x candi_size
Return:
score: B x S x candi_size
'''
hidden = hidden.unsqueeze(2).repeat(1, 1, candi_embeddings.size(1), 1) # B x S x candi_size x H
candi_embeddings = candi_embeddings.unsqueeze(1).repeat(1, hidden.size(1), 1, 1) # B x S x candi_size x H
candi_mask = candi_mask.unsqueeze(1).repeat(1, hidden.size(1), 1) # B x S x candi_size
energy_in = torch.cat((hidden, candi_embeddings), -1) # B x S x candi_size x 2H
score = self.score(torch.tanh(self.attn(energy_in))).squeeze(-1) # B x S x candi_size
if candi_mask is not None:
score = score.masked_fill_(~candi_mask, -1e12)
return score |