-
Notifications
You must be signed in to change notification settings - Fork 186
/
decode_full_model.py
169 lines (147 loc) · 6.2 KB
/
decode_full_model.py
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
""" run decoding of rnn-ext + abs + RL (+ rerank)"""
import argparse
import json
import os
from os.path import join
from datetime import timedelta
from time import time
from collections import Counter, defaultdict
from itertools import product
from functools import reduce
import operator as op
from cytoolz import identity, concat, curry
import torch
from torch.utils.data import DataLoader
from torch import multiprocessing as mp
from data.batcher import tokenize
from decoding import Abstractor, RLExtractor, DecodeDataset, BeamAbstractor
from decoding import make_html_safe
def decode(save_path, model_dir, split, batch_size,
beam_size, diverse, max_len, cuda):
start = time()
# setup model
with open(join(model_dir, 'meta.json')) as f:
meta = json.loads(f.read())
if meta['net_args']['abstractor'] is None:
# NOTE: if no abstractor is provided then
# the whole model would be extractive summarization
assert beam_size == 1
abstractor = identity
else:
if beam_size == 1:
abstractor = Abstractor(join(model_dir, 'abstractor'),
max_len, cuda)
else:
abstractor = BeamAbstractor(join(model_dir, 'abstractor'),
max_len, cuda)
extractor = RLExtractor(model_dir, cuda=cuda)
# setup loader
def coll(batch):
articles = list(filter(bool, batch))
return articles
dataset = DecodeDataset(split)
n_data = len(dataset)
loader = DataLoader(
dataset, batch_size=batch_size, shuffle=False, num_workers=4,
collate_fn=coll
)
# prepare save paths and logs
os.makedirs(join(save_path, 'output'))
dec_log = {}
dec_log['abstractor'] = meta['net_args']['abstractor']
dec_log['extractor'] = meta['net_args']['extractor']
dec_log['rl'] = True
dec_log['split'] = split
dec_log['beam'] = beam_size
dec_log['diverse'] = diverse
with open(join(save_path, 'log.json'), 'w') as f:
json.dump(dec_log, f, indent=4)
# Decoding
i = 0
with torch.no_grad():
for i_debug, raw_article_batch in enumerate(loader):
tokenized_article_batch = map(tokenize(None), raw_article_batch)
ext_arts = []
ext_inds = []
for raw_art_sents in tokenized_article_batch:
ext = extractor(raw_art_sents)[:-1] # exclude EOE
if not ext:
# use top-5 if nothing is extracted
# in some rare cases rnn-ext does not extract at all
ext = list(range(5))[:len(raw_art_sents)]
else:
ext = [i.item() for i in ext]
ext_inds += [(len(ext_arts), len(ext))]
ext_arts += [raw_art_sents[i] for i in ext]
if beam_size > 1:
all_beams = abstractor(ext_arts, beam_size, diverse)
dec_outs = rerank_mp(all_beams, ext_inds)
else:
dec_outs = abstractor(ext_arts)
assert i == batch_size*i_debug
for j, n in ext_inds:
decoded_sents = [' '.join(dec) for dec in dec_outs[j:j+n]]
with open(join(save_path, 'output/{}.dec'.format(i)),
'w') as f:
f.write(make_html_safe('\n'.join(decoded_sents)))
i += 1
print('{}/{} ({:.2f}%) decoded in {} seconds\r'.format(
i, n_data, i/n_data*100,
timedelta(seconds=int(time()-start))
), end='')
print()
_PRUNE = defaultdict(
lambda: 2,
{1:5, 2:5, 3:5, 4:5, 5:5, 6:4, 7:3, 8:3}
)
def rerank(all_beams, ext_inds):
beam_lists = (all_beams[i: i+n] for i, n in ext_inds if n > 0)
return list(concat(map(rerank_one, beam_lists)))
def rerank_mp(all_beams, ext_inds):
beam_lists = [all_beams[i: i+n] for i, n in ext_inds if n > 0]
with mp.Pool(8) as pool:
reranked = pool.map(rerank_one, beam_lists)
return list(concat(reranked))
def rerank_one(beams):
@curry
def process_beam(beam, n):
for b in beam[:n]:
b.gram_cnt = Counter(_make_n_gram(b.sequence))
return beam[:n]
beams = map(process_beam(n=_PRUNE[len(beams)]), beams)
best_hyps = max(product(*beams), key=_compute_score)
dec_outs = [h.sequence for h in best_hyps]
return dec_outs
def _make_n_gram(sequence, n=2):
return (tuple(sequence[i:i+n]) for i in range(len(sequence)-(n-1)))
def _compute_score(hyps):
all_cnt = reduce(op.iadd, (h.gram_cnt for h in hyps), Counter())
repeat = sum(c-1 for g, c in all_cnt.items() if c > 1)
lp = sum(h.logprob for h in hyps) / sum(len(h.sequence) for h in hyps)
return (-repeat, lp)
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='run decoding of the full model (RL)')
parser.add_argument('--path', required=True, help='path to store/eval')
parser.add_argument('--model_dir', help='root of the full model')
# dataset split
data = parser.add_mutually_exclusive_group(required=True)
data.add_argument('--val', action='store_true', help='use validation set')
data.add_argument('--test', action='store_true', help='use test set')
# decode options
parser.add_argument('--batch', type=int, action='store', default=32,
help='batch size of faster decoding')
parser.add_argument('--beam', type=int, action='store', default=1,
help='beam size for beam-search (reranking included)')
parser.add_argument('--div', type=float, action='store', default=1.0,
help='diverse ratio for the diverse beam-search')
parser.add_argument('--max_dec_word', type=int, action='store', default=30,
help='maximun words to be decoded for the abstractor')
parser.add_argument('--no-cuda', action='store_true',
help='disable GPU training')
args = parser.parse_args()
args.cuda = torch.cuda.is_available() and not args.no_cuda
data_split = 'test' if args.test else 'val'
decode(args.path, args.model_dir,
data_split, args.batch, args.beam, args.div,
args.max_dec_word, args.cuda)