-
Notifications
You must be signed in to change notification settings - Fork 19
/
decode_full_model.py
180 lines (152 loc) · 6.74 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
170
171
172
173
174
175
176
177
178
179
180
""" run decoding of rnn-ext + abs + RL (+ rerank)"""
import argparse
import json
import os
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 data.data import CnnDmDataset
from utils.decoding import Abstractor, RLExtractor, BeamAbstractor, make_html_safe
class DecodeDataset(CnnDmDataset):
""" get the article sentences only (for decoding use)"""
def __init__(self, split, data_dir):
assert split in ['val', 'test']
super().__init__(split, data_dir)
def __getitem__(self, i):
js_data = super().__getitem__(i)
art_sents = js_data['article']
return art_sents
def decode(data_dir, save_path, model_dir, split, batch_size,
beam_size, diverse, max_len, cuda):
start = time()
# setup model
with open(os.path.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(os.path.join(model_dir, 'abstractor'),
max_len, cuda)
else:
abstractor = BeamAbstractor(os.path.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, data_dir)
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(os.path.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(os.path.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
ext = extractor(raw_art_sents)
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)]
ext = list(range(1))[:len(raw_art_sents)]
print("none")
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(os.path.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='')
_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)+1)
return (-repeat, lp)
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='run decoding of the full model (RL)')
parser.add_argument('--data', required=True, help='data directories')
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=5,
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.data, args.path, args.model_dir, data_split,
args.batch, args.beam, args.div,
args.max_dec_word, args.cuda)