-
Notifications
You must be signed in to change notification settings - Fork 0
/
sample.py
50 lines (43 loc) · 1.91 KB
/
sample.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
from __future__ import print_function
import numpy as np
import tensorflow as tf
import argparse
import time
import os
from six.moves import cPickle
from utils import TextLoader
from .model import Model
def main():
#parses arguments from command line
parser = argparse.ArgumentParser()
parser.add_argument('--save_dir', type=str, default='save',
help='model directory to store checkpointed models')
parser.add_argument('-n', type=int, default=200,
help='number of words to sample')
parser.add_argument('--prime', type=str, default=' ',
help='prime text')
parser.add_argument('--pick', type=int, default=1,
help='1 = weighted pick, 2 = beam search pick')
parser.add_argument('--sample', type=int, default=1,
help='0 to use max at each timestep, 1 to sample at each timestep, 2 to sample on spaces')
args = parser.parse_args()
sample(args) #calls sample function
def sample(args):
#open config and words_vocab
with open(os.path.join(args.save_dir, 'config.pkl'), 'rb') as f:
saved_args = cPickle.load(f)
with open(os.path.join(args.save_dir, 'words_vocab.pkl'), 'rb') as f:
words, vocab = cPickle.load(f)
#load saved model
model = Model(saved_args, True)
with tf.Session() as sess:
tf.global_variables_initializer().run()
saver = tf.train.Saver(tf.global_variables())
ckpt = tf.train.get_checkpoint_state(args.save_dir)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path) #restore session with model checkpoint
#print sample from given model
print(model.sample(sess, words, vocab, args.n, args.prime, args.sample, args.pick))
return model.sample(sess, words, vocab, args.n, args.prime, args.sample, args.pick)
if __name__ == '__main__':
main()