-
Notifications
You must be signed in to change notification settings - Fork 1
/
auto_reply_msg.py
74 lines (62 loc) · 1.95 KB
/
auto_reply_msg.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
#!/usr/bin/env python3
import fire
import json
import os
import numpy as np
import tensorflow as tf
import model, sample, encoder
def interact_model(
message="A quick look at the",
model_name='345M',
seed=5,
nsamples=1,
batch_size=1,
length=50,
temperature=0.9,
top_k=20,
top_p=0.9,
):
if batch_size is None:
batch_size = 1
assert nsamples % batch_size == 0
enc = encoder.get_encoder(model_name)
hparams = model.default_hparams()
path = os.path.dirname(__file__)
with open(os.path.join(path, 'models',
model_name,
'hparams.json')) as f:
hparams.override_from_dict(json.load(f))
if length is None:
length = hparams.n_ctx // 2
elif length > hparams.n_ctx:
raise ValueError("Can't get samples longer"
" than window size: %s" % hparams.n_ctx)
with tf.Session(graph=tf.Graph()) as sess:
context = tf.placeholder(tf.int32, [batch_size, None])
np.random.seed(seed)
tf.set_random_seed(seed)
output = sample.sample_sequence(
hparams=hparams, length=length,
context=context,
batch_size=batch_size,
temperature=temperature,
top_k=top_k,
top_p=top_p
)
saver = tf.train.Saver()
ckpt = tf.train.latest_checkpoint(os.path.join(path, 'models',
model_name))
saver.restore(sess, ckpt)
if message == "":
return -1
raw_text = message
context_tokens = enc.encode(raw_text)
out = sess.run(output, feed_dict={
context: [context_tokens for _ in range(batch_size)]
})[:, len(context_tokens):]
text = []
for i in range(batch_size):
text.append(enc.decode(out[i]))
return text
if __name__ == '__main__':
print(fire.Fire(interact_model))