forked from kelvinxu/arctic-captions
-
Notifications
You must be signed in to change notification settings - Fork 0
/
generate_caps.py
133 lines (115 loc) · 4.79 KB
/
generate_caps.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
"""
Sampling script for attention models
Works on CPU with support for multi-process
"""
import argparse
import numpy
import cPickle as pkl
import json
from progress.bar import Bar
from capgen import build_sampler, gen_sample, \
load_params, \
init_params, \
init_tparams, \
get_dataset \
import theano
from theano import tensor
from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
from IPython import embed
def create_sample(tparams, f_init, f_next, context, options, trng, k, normalize):
sample, score = gen_sample(tparams, f_init, f_next, context, options,
trng=trng, k=k, maxlen=200, stochastic=False)
# adjust for length bias
if normalize:
lengths = numpy.array([len(s) for s in sample])
score = score / lengths
sidx = numpy.argmin(score)
return sample[sidx]
def main(model, saveto, k=1, normalize=False, zero_pad=False, datasets='dev,test', data_path='./', sampling=False, pkl_name=None):
# load model model_options
if pkl_name is None:
pkl_name = model
with open('%s.pkl'% pkl_name, 'rb') as f:
options = pkl.load(f)
# fetch data, skip ones we aren't using to save time
load_data, prepare_data = get_dataset(options['dataset'])
_, valid, test, worddict = load_data(load_train=False, load_dev=True if 'dev' in datasets else False,
load_test=True if 'test' in datasets else False, path=data_path)
# <eos> means end of sequence (aka periods), UNK means unknown
word_idict = dict()
for kk, vv in worddict.iteritems():
word_idict[vv] = kk
word_idict[0] = '<eos>'
word_idict[1] = 'UNK'
# build sampler
trng = RandomStreams(1234)
# this is zero indicate we are not using dropout in the graph
use_noise = theano.shared(numpy.float32(0.), name='use_noise')
# get the parameters
params = init_params(options)
params = load_params(model, params)
tparams = init_tparams(params)
# build the sampling computational graph
# see capgen.py for more detailed explanations
f_init, f_next = build_sampler(tparams, options, use_noise, trng, sampling=sampling)
# index -> words
def _seqs2words(cc):
ww = []
for w in cc:
if w == 0:
break
ww.append(word_idict[w])
return ' '.join(ww)
# unsparsify, reshape, and queue
def _send_job(context):
cc = context.todense().reshape([14*14,512])
if zero_pad:
cc0 = numpy.zeros((cc.shape[0]+1, cc.shape[1])).astype('float32')
cc0[:-1,:] = cc
else:
cc0 = cc
return create_sample(tparams, f_init, f_next, cc0, options, trng, k, normalize)
ds = datasets.strip().split(',')
# send all the features for the various datasets
for dd in ds:
if dd == 'dev':
bar = Bar('Development Set...', max=len(valid[1]))
caps = []
for i in range(len(valid[1])):
sample = _send_job(valid[1][i])
cap = _seqs2words(sample)
caps.append(cap)
with open(saveto+'_status.json', 'w') as f:
json.dump({'current': i, 'total': len(valid[1])}, f)
bar.next()
bar.finish()
with open(saveto, 'w') as f:
print >>f, '\n'.join(caps)
print 'Done'
if dd == 'test':
print 'Test Set...',
caps = []
for i in range(len(test[1])):
sample = _send_job(test[1][i])
cap = _seqs2words(sample)
caps.append(cap)
with open(saveto+'_status.json', 'w') as f:
json.dump({'current': i, 'total': len(test[1])}, f)
bar.next()
bar.finish()
with open(saveto, 'w') as f:
print >>f, '\n'.join(caps)
print 'Done'
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-k', type=int, default=1)
parser.add_argument('-sampling', action="store_true", default=False) # this only matters for hard attention
parser.add_argument('-n', action="store_true", default=False)
parser.add_argument('-z', action="store_true", default=False)
parser.add_argument('-d', type=str, default='dev')
parser.add_argument('-pkl_name', type=str, default=None, help="name of pickle file (without the .pkl)")
parser.add_argument('data', type=str)
parser.add_argument('model', type=str)
parser.add_argument('saveto', type=str)
args = parser.parse_args()
main(args.model, args.saveto, k=args.k, zero_pad=args.z, pkl_name=args.pkl_name, data_path=args.data, normalize=args.n, datasets=args.d, sampling=args.sampling)