wehub-resource-sync
297 lines
13 KiB
Python
297 lines
13 KiB
Python
import torch
|
|
from modules import BiLSTM, GraphTrans, MSA
|
|
from torch import nn
|
|
from utlis import *
|
|
|
|
import dgl
|
|
|
|
|
|
class GraphWriter(nn.Module):
|
|
def __init__(self, args):
|
|
super(GraphWriter, self).__init__()
|
|
self.args = args
|
|
if args.title:
|
|
self.title_emb = nn.Embedding(
|
|
len(args.title_vocab), args.nhid, padding_idx=0
|
|
)
|
|
self.title_enc = BiLSTM(args, enc_type="title")
|
|
self.title_attn = MSA(args)
|
|
self.ent_emb = nn.Embedding(
|
|
len(args.ent_text_vocab), args.nhid, padding_idx=0
|
|
)
|
|
self.tar_emb = nn.Embedding(
|
|
len(args.text_vocab), args.nhid, padding_idx=0
|
|
)
|
|
if args.title:
|
|
nn.init.xavier_normal_(self.title_emb.weight)
|
|
nn.init.xavier_normal_(self.ent_emb.weight)
|
|
self.rel_emb = nn.Embedding(
|
|
len(args.rel_vocab), args.nhid, padding_idx=0
|
|
)
|
|
nn.init.xavier_normal_(self.rel_emb.weight)
|
|
self.decode_lstm = nn.LSTMCell(args.dec_ninp, args.nhid)
|
|
self.ent_enc = BiLSTM(args, enc_type="entity")
|
|
self.graph_enc = GraphTrans(args)
|
|
self.ent_attn = MSA(args)
|
|
self.copy_attn = MSA(args, mode="copy")
|
|
self.copy_fc = nn.Linear(args.dec_ninp, 1)
|
|
self.pred_v_fc = nn.Linear(args.dec_ninp, len(args.text_vocab))
|
|
|
|
def enc_forward(
|
|
self, batch, ent_mask, ent_text_mask, ent_len, rel_mask, title_mask
|
|
):
|
|
title_enc = None
|
|
if self.args.title:
|
|
title_enc = self.title_enc(
|
|
self.title_emb(batch["title"]), title_mask
|
|
)
|
|
ent_enc = self.ent_enc(
|
|
self.ent_emb(batch["ent_text"]),
|
|
ent_text_mask,
|
|
ent_len=batch["ent_len"],
|
|
)
|
|
rel_emb = self.rel_emb(batch["rel"])
|
|
g_ent, g_root = self.graph_enc(
|
|
ent_enc, ent_mask, ent_len, rel_emb, rel_mask, batch["graph"]
|
|
)
|
|
return g_ent, g_root, title_enc, ent_enc
|
|
|
|
def forward(self, batch, beam_size=-1):
|
|
ent_mask = len2mask(batch["ent_len"], self.args.device)
|
|
ent_text_mask = batch["ent_text"] == 0
|
|
rel_mask = batch["rel"] == 0 # 0 means the <PAD>
|
|
title_mask = batch["title"] == 0
|
|
g_ent, g_root, title_enc, ent_enc = self.enc_forward(
|
|
batch,
|
|
ent_mask,
|
|
ent_text_mask,
|
|
batch["ent_len"],
|
|
rel_mask,
|
|
title_mask,
|
|
)
|
|
|
|
_h, _c = g_root, g_root.clone().detach()
|
|
ctx = _h + self.ent_attn(_h, g_ent, mask=ent_mask)
|
|
if self.args.title:
|
|
attn = _h + self.title_attn(_h, title_enc, mask=title_mask)
|
|
ctx = torch.cat([ctx, attn], 1)
|
|
if beam_size < 1:
|
|
# training
|
|
outs = []
|
|
tar_inp = self.tar_emb(batch["text"].transpose(0, 1))
|
|
for t, xt in enumerate(tar_inp):
|
|
_xt = torch.cat([ctx, xt], 1)
|
|
_h, _c = self.decode_lstm(_xt, (_h, _c))
|
|
ctx = _h + self.ent_attn(_h, g_ent, mask=ent_mask)
|
|
if self.args.title:
|
|
attn = _h + self.title_attn(_h, title_enc, mask=title_mask)
|
|
ctx = torch.cat([ctx, attn], 1)
|
|
outs.append(torch.cat([_h, ctx], 1))
|
|
outs = torch.stack(outs, 1)
|
|
copy_gate = torch.sigmoid(self.copy_fc(outs))
|
|
EPSI = 1e-6
|
|
# copy
|
|
pred_v = torch.log(copy_gate + EPSI) + torch.log_softmax(
|
|
self.pred_v_fc(outs), -1
|
|
)
|
|
pred_c = torch.log((1.0 - copy_gate) + EPSI) + torch.log_softmax(
|
|
self.copy_attn(outs, ent_enc, mask=ent_mask), -1
|
|
)
|
|
pred = torch.cat([pred_v, pred_c], -1)
|
|
return pred
|
|
else:
|
|
if beam_size == 1:
|
|
# greedy
|
|
device = g_ent.device
|
|
B = g_ent.shape[0]
|
|
ent_type = batch["ent_type"].view(B, -1)
|
|
seq = (
|
|
torch.ones(
|
|
B,
|
|
)
|
|
.long()
|
|
.to(device)
|
|
* self.args.text_vocab("<BOS>")
|
|
).unsqueeze(1)
|
|
for t in range(self.args.beam_max_len):
|
|
_inp = replace_ent(
|
|
seq[:, -1], ent_type, len(self.args.text_vocab)
|
|
)
|
|
xt = self.tar_emb(_inp)
|
|
_xt = torch.cat([ctx, xt], 1)
|
|
_h, _c = self.decode_lstm(_xt, (_h, _c))
|
|
ctx = _h + self.ent_attn(_h, g_ent, mask=ent_mask)
|
|
if self.args.title:
|
|
attn = _h + self.title_attn(
|
|
_h, title_enc, mask=title_mask
|
|
)
|
|
ctx = torch.cat([ctx, attn], 1)
|
|
_y = torch.cat([_h, ctx], 1)
|
|
copy_gate = torch.sigmoid(self.copy_fc(_y))
|
|
pred_v = torch.log(copy_gate) + torch.log_softmax(
|
|
self.pred_v_fc(_y), -1
|
|
)
|
|
pred_c = torch.log((1.0 - copy_gate)) + torch.log_softmax(
|
|
self.copy_attn(
|
|
_y.unsqueeze(1), ent_enc, mask=ent_mask
|
|
).squeeze(1),
|
|
-1,
|
|
)
|
|
pred = torch.cat([pred_v, pred_c], -1).view(B, -1)
|
|
for ban_item in ["<BOS>", "<PAD>", "<UNK>"]:
|
|
pred[:, self.args.text_vocab(ban_item)] = -1e8
|
|
_, word = pred.max(-1)
|
|
seq = torch.cat([seq, word.unsqueeze(1)], 1)
|
|
return seq
|
|
else:
|
|
# beam search
|
|
device = g_ent.device
|
|
B = g_ent.shape[0]
|
|
BSZ = B * beam_size
|
|
_h = _h.view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
|
|
_c = _c.view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
|
|
ent_mask = (
|
|
ent_mask.view(B, 1, -1)
|
|
.repeat(1, beam_size, 1)
|
|
.view(BSZ, -1)
|
|
)
|
|
if self.args.title:
|
|
title_mask = (
|
|
title_mask.view(B, 1, -1)
|
|
.repeat(1, beam_size, 1)
|
|
.view(BSZ, -1)
|
|
)
|
|
title_enc = (
|
|
title_enc.view(B, 1, title_enc.size(1), -1)
|
|
.repeat(1, beam_size, 1, 1)
|
|
.view(BSZ, title_enc.size(1), -1)
|
|
)
|
|
ctx = ctx.view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
|
|
ent_type = (
|
|
batch["ent_type"]
|
|
.view(B, 1, -1)
|
|
.repeat(1, beam_size, 1)
|
|
.view(BSZ, -1)
|
|
)
|
|
g_ent = (
|
|
g_ent.view(B, 1, g_ent.size(1), -1)
|
|
.repeat(1, beam_size, 1, 1)
|
|
.view(BSZ, g_ent.size(1), -1)
|
|
)
|
|
ent_enc = (
|
|
ent_enc.view(B, 1, ent_enc.size(1), -1)
|
|
.repeat(1, beam_size, 1, 1)
|
|
.view(BSZ, ent_enc.size(1), -1)
|
|
)
|
|
|
|
beam_best = torch.zeros(B).to(device) - 1e9
|
|
beam_best_seq = [None] * B
|
|
beam_seq = (
|
|
torch.ones(B, beam_size).long().to(device)
|
|
* self.args.text_vocab("<BOS>")
|
|
).unsqueeze(-1)
|
|
beam_score = torch.zeros(B, beam_size).to(device)
|
|
done_flag = torch.zeros(B, beam_size)
|
|
for t in range(self.args.beam_max_len):
|
|
_inp = replace_ent(
|
|
beam_seq[:, :, -1].view(-1),
|
|
ent_type,
|
|
len(self.args.text_vocab),
|
|
)
|
|
xt = self.tar_emb(_inp)
|
|
_xt = torch.cat([ctx, xt], 1)
|
|
_h, _c = self.decode_lstm(_xt, (_h, _c))
|
|
ctx = _h + self.ent_attn(_h, g_ent, mask=ent_mask)
|
|
if self.args.title:
|
|
attn = _h + self.title_attn(
|
|
_h, title_enc, mask=title_mask
|
|
)
|
|
ctx = torch.cat([ctx, attn], 1)
|
|
_y = torch.cat([_h, ctx], 1)
|
|
copy_gate = torch.sigmoid(self.copy_fc(_y))
|
|
pred_v = torch.log(copy_gate) + torch.log_softmax(
|
|
self.pred_v_fc(_y), -1
|
|
)
|
|
pred_c = torch.log((1.0 - copy_gate)) + torch.log_softmax(
|
|
self.copy_attn(
|
|
_y.unsqueeze(1), ent_enc, mask=ent_mask
|
|
).squeeze(1),
|
|
-1,
|
|
)
|
|
pred = torch.cat([pred_v, pred_c], -1).view(
|
|
B, beam_size, -1
|
|
)
|
|
for ban_item in ["<BOS>", "<PAD>", "<UNK>"]:
|
|
pred[:, :, self.args.text_vocab(ban_item)] = -1e8
|
|
if t == self.args.beam_max_len - 1: # force ending
|
|
tt = pred[:, :, self.args.text_vocab("<EOS>")]
|
|
pred = pred * 0 - 1e8
|
|
pred[:, :, self.args.text_vocab("<EOS>")] = tt
|
|
cum_score = beam_score.view(B, beam_size, 1) + pred
|
|
score, word = cum_score.topk(
|
|
dim=-1, k=beam_size
|
|
) # B, beam_size, beam_size
|
|
score, word = score.view(B, -1), word.view(B, -1)
|
|
eos_idx = self.args.text_vocab("<EOS>")
|
|
if beam_seq.size(2) == 1:
|
|
new_idx = torch.arange(beam_size).to(word)
|
|
new_idx = new_idx[None, :].repeat(B, 1)
|
|
else:
|
|
_, new_idx = score.topk(dim=-1, k=beam_size)
|
|
new_src, new_score, new_word, new_done = [], [], [], []
|
|
LP = beam_seq.size(2) ** self.args.lp
|
|
for i in range(B):
|
|
for j in range(beam_size):
|
|
tmp_score = score[i][new_idx[i][j]]
|
|
tmp_word = word[i][new_idx[i][j]]
|
|
src_idx = new_idx[i][j] // beam_size
|
|
new_src.append(src_idx)
|
|
if tmp_word == eos_idx:
|
|
new_score.append(-1e8)
|
|
else:
|
|
new_score.append(tmp_score)
|
|
new_word.append(tmp_word)
|
|
|
|
if (
|
|
tmp_word == eos_idx
|
|
and done_flag[i][src_idx] == 0
|
|
and tmp_score / LP > beam_best[i]
|
|
):
|
|
beam_best[i] = tmp_score / LP
|
|
beam_best_seq[i] = beam_seq[i][src_idx]
|
|
if tmp_word == eos_idx:
|
|
new_done.append(1)
|
|
else:
|
|
new_done.append(done_flag[i][src_idx])
|
|
new_score = (
|
|
torch.Tensor(new_score)
|
|
.view(B, beam_size)
|
|
.to(beam_score)
|
|
)
|
|
new_word = (
|
|
torch.Tensor(new_word).view(B, beam_size).to(beam_seq)
|
|
)
|
|
new_src = (
|
|
torch.LongTensor(new_src).view(B, beam_size).to(device)
|
|
)
|
|
new_done = (
|
|
torch.Tensor(new_done).view(B, beam_size).to(done_flag)
|
|
)
|
|
beam_score = new_score
|
|
done_flag = new_done
|
|
beam_seq = beam_seq.view(B, beam_size, -1)[
|
|
torch.arange(B)[:, None].to(device), new_src
|
|
]
|
|
beam_seq = torch.cat([beam_seq, new_word.unsqueeze(2)], 2)
|
|
_h = _h.view(B, beam_size, -1)[
|
|
torch.arange(B)[:, None].to(device), new_src
|
|
].view(BSZ, -1)
|
|
_c = _c.view(B, beam_size, -1)[
|
|
torch.arange(B)[:, None].to(device), new_src
|
|
].view(BSZ, -1)
|
|
ctx = ctx.view(B, beam_size, -1)[
|
|
torch.arange(B)[:, None].to(device), new_src
|
|
].view(BSZ, -1)
|
|
|
|
return beam_best_seq
|