def load_ext_net(ext_dir):
ext_meta = json.load(open(join(ext_dir, 'meta.json')))
assert ext_meta['net'] == 'ml_rnn_extractor'
ext_ckpt = load_best_ckpt(ext_dir)
ext_args = ext_meta['net_args']
vocab = pkl.load(open(join(ext_dir, 'vocab.pkl'), 'rb'))
ext = PtrExtractSumm(**ext_args)
ext.load_state_dict(ext_ckpt)
return ext, vocab
def load_ext_net(ext_dir):
ext_meta = json.load(open(join(ext_dir, 'meta.json')))
assert ext_meta['net'] == 'ml_rnn_extractor' or ext_meta[
'net'] == "ml_entity_extractor"
ext_ckpt = load_best_ckpt(ext_dir)
ext_args = ext_meta['net_args']
vocab = pkl.load(open(join(ext_dir, 'vocab.pkl'), 'rb'))
if ext_meta['net'] == 'ml_rnn_extractor':
ext = PtrExtractSumm(**ext_args)
elif ext_meta['net'] == "ml_entity_extractor":
ext = PtrExtractSummEntity(**ext_args)
else:
raise Exception('not implemented')
ext.load_state_dict(ext_ckpt)
return ext, vocab
def load_ext_net(ext_dir):
ext_meta = json.load(open(join(ext_dir, 'meta.json')))
assert ext_meta['net'] in [
'ml_rnn_extractor', "ml_gat_extractor", "ml_subgraph_gat_extractor"
]
net_name = ext_meta['net']
ext_ckpt = load_best_ckpt(ext_dir)
ext_args = ext_meta['net_args']
vocab = pkl.load(open(join(ext_dir, 'vocab.pkl'), 'rb'))
if ext_meta['net'] == 'ml_rnn_extractor':
ext = PtrExtractSumm(**ext_args)
elif ext_meta['net'] == "ml_gat_extractor":
ext = PtrExtractSummGAT(**ext_args)
elif ext_meta['net'] == "ml_subgraph_gat_extractor":
ext = PtrExtractSummSubgraph(**ext_args)
else:
raise Exception('not implemented')
ext.load_state_dict(ext_ckpt)
return ext, vocab
def load_dis_net(emb_dim,
lstm_hidden,
lstm_layer,
bert_config,
dis_pretrain_file,
load=True,
cuda=True):
dis = PtrExtractSumm(emb_dim=emb_dim,
lstm_hidden=lstm_hidden,
lstm_layer=lstm_layer,
bert_config=bert_config)
dis = PolicyGradient(dis.transformer, dis._extractor)
if load:
print("Restoring all non-adagrad variables from {}...".format(
dis_pretrain_file))
state_dict = torch.load(dis_pretrain_file)['state_dict']
dis.load_state_dict(state_dict)
if cuda:
dis = dis.cuda()
return dis
