def setup(opts):
initial_context = tf.placeholder(tf.int32, [batch_size_per_chunk, None])
p_for_topp = tf.placeholder(tf.float32, [batch_size_per_chunk])
eos_token = tf.placeholder(tf.int32, [])
tokens, probs = sample(news_config=news_config,
initial_context=initial_context,
eos_token=eos_token,
ignore_ids=None,
p_for_topp=p_for_topp,
do_topk=False)
saver = tf.train.Saver()
checkpoint_folder = glob_dir(opts['checkpoint_dir'])[0]
checkpoint_path = '.'.join(glob_dir(checkpoint_folder)[0].split('.')[:-1])
saver.restore(sess, checkpoint_path)
return {
'tokens': tokens,
'probs': probs,
'initial_context': initial_context,
'eos_token': eos_token,
'p_for_topp': p_for_topp
}
args.batch_size, max_batch_size, num_chunks, batch_size_per_chunk), flush=True)
# This controls the top p for each generation.
top_p = np.ones((num_chunks, batch_size_per_chunk), dtype=np.float32) * args.top_p
# with open(args.metadata_fn, 'r') as f:
# articles = [json.loads(l) for i, l in enumerate(f) if i % args.num_folds == args.fold]
tf_config = tf.ConfigProto(allow_soft_placement=True)
with tf.Session(config=tf_config, graph=tf.Graph()) as sess:
initial_context = tf.placeholder(tf.int32, [batch_size_per_chunk, None])
p_for_topp = tf.placeholder(tf.float32, [batch_size_per_chunk])
eos_token = tf.placeholder(tf.int32, [])
tokens, probs = sample(news_config=news_config, initial_context=initial_context,
eos_token=eos_token, ignore_ids=None, p_for_topp=p_for_topp,
do_topk=False)
saver = tf.train.Saver()
saver.restore(sess, args.model_ckpt)
print('Loaded model.')
text = input()
while text != "":
for i in range(args.samples):
print("Sample,", i + 1, " of ", args.samples)
# Let's go!
encoded = _tokenize_article_pieces(encoder, text)
context_formatted = []
# for key in ['domain', 'date', 'authors', 'title', 'article']:
# if key != args.target:
# context_formatted.extend(article_pieces.pop(key, []))