def __init__(self, ext_dir, cuda=True):
ext_meta = json.load(open(join(ext_dir, 'meta.json')))
assert ext_meta['net'] == 'rnn-ext_abs_rl'
ext_args = ext_meta['net_args']['extractor']['net_args']
word2id = pkl.load(open(join(ext_dir, 'agent_vocab.pkl'), 'rb'))
extractor = PtrExtractSumm(**ext_args)
agent = ActorCritic(extractor, ArticleBatcher(word2id, cuda))
ext_ckpt = load_best_ckpt(ext_dir, reverse=True)
agent.load_state_dict(ext_ckpt)
self._device = torch.device('cuda' if cuda else 'cpu')
self._net = agent.to(self._device)
self._word2id = word2id
self._id2word = {i: w for w, i in word2id.items()}
def load_rl_ckpt(abs_dir, ext_dir, cuda):
ext_meta = json.load(open(join(ext_dir, 'meta.json')))
assert ext_meta['net'] == 'rnn-ext_abs_rl'
ext_args = ext_meta['net_args']['extractor']['net_args']
vocab = pkl.load(open(join(ext_dir, 'agent_vocab.pkl'), 'rb'))
extractor = PtrExtractSumm(**ext_args)
abstractor_sent = Abstractor(abs_dir, MAX_ABS_LEN, cuda)
agent = ActorCritic(extractor._sent_enc, extractor._art_enc,
extractor._extractor, ArticleBatcher(vocab, cuda))
target_agent = ActorCritic(extractor._sent_enc,
extractor._art_enc, extractor._extractor,
ArticleBatcher(vocab, cuda))
ext_ckpt = load_best_ckpt(ext_dir, reverse=True)
agent.load_state_dict(ext_ckpt)
device = torch.device('cuda' if cuda else 'cpu')
agent = agent.to(device)
target_agent = target_agent.to(device)
return agent, target_agent, vocab, abstractor_sent, ext_meta
def __init__(self, ext_dir, cuda=True):
ext_meta = json.load(open(join(ext_dir, 'meta.json')))
assert ext_meta['net'] == 'rnn-ext_abs_rl'
ext_args = ext_meta['net_args']['extractor']['net_args']
word2id = pkl.load(open(join(ext_dir, 'agent_vocab.pkl'), 'rb'))
extractor = PtrExtractSumm(**ext_args)
agent = ActorCritic(extractor._sent_enc,
extractor._art_enc,
extractor._extractor,
ArticleBatcher(word2id, cuda))
ext_ckpt = load_best_ckpt(ext_dir, reverse=True)
agent.load_state_dict(ext_ckpt)
self._device = torch.device('cuda' if cuda else 'cpu')
self._net = agent.to(self._device)
self._word2id = word2id
self._id2word = {i: w for w, i in word2id.items()}
