mirror of
https://github.com/wassname/Clover-Edition.git
synced 2026-09-20 12:20:53 +08:00
143 lines
3.4 KiB
Python
143 lines
3.4 KiB
Python
import json
|
|
import os
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
from src.model import *
|
|
from tensorflow.contrib import predictor
|
|
from src.sample import *
|
|
from src.encoder import *
|
|
import pdb
|
|
|
|
pos_action_starts = ["You attack", "You tell", "You use", "You go"]
|
|
|
|
|
|
class TFGenerator():
|
|
|
|
def __init__(self, sess, length=75, temperature=0.9, top_k=40):
|
|
|
|
seed = None
|
|
batch_size=1
|
|
model_path='gpt2/models/117M'
|
|
self.sess = sess
|
|
|
|
self.enc = encoder.get_encoder(model_path)
|
|
hparams = model.default_hparams()
|
|
with open(os.path.join(model_path, 'hparams.json')) as f:
|
|
hparams.override_from_dict(json.load(f))
|
|
|
|
pdb.set_trace()
|
|
|
|
self.context = tf.placeholder(tf.int32, [batch_size, None])
|
|
np.random.seed(seed)
|
|
tf.set_random_seed(seed)
|
|
self.output = sample.sample_sequence(
|
|
hparams=hparams, length=length,
|
|
context=self.context,
|
|
batch_size=batch_size,
|
|
)
|
|
|
|
saver = tf.train.Saver()
|
|
ckpt = tf.train.latest_checkpoint(model_path)
|
|
saver.restore(self.sess, ckpt)
|
|
|
|
def generate(self, prompt, options={}):
|
|
context_tokens = self.enc.encode(prompt)
|
|
out = self.sess.run(self.output, feed_dict={
|
|
self.context: [context_tokens for _ in range(1)]
|
|
})[:, len(context_tokens):]
|
|
|
|
text = self.enc.decode(out[0])
|
|
return text
|
|
|
|
def save_model():
|
|
length=75
|
|
temperature=0.9
|
|
top_k=40
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
|
|
|
|
with tf.Session() as sess:
|
|
seed = None
|
|
batch_size=None
|
|
model_path='models/774M'
|
|
|
|
hparams = default_hparams()
|
|
with open(os.path.join(model_path, 'hparams.json')) as f:
|
|
hparams.override_from_dict(json.load(f))
|
|
|
|
context = tf.placeholder(tf.int32, [batch_size, None])
|
|
np.random.seed(seed)
|
|
tf.set_random_seed(seed)
|
|
output = sample_sequence(
|
|
hparams=hparams, length=length,
|
|
context=context,
|
|
batch_size=batch_size,
|
|
)
|
|
print("***********************",type(output))
|
|
|
|
saver = tf.train.Saver()
|
|
ckpt = tf.train.latest_checkpoint(model_path)
|
|
saver.restore(sess, ckpt)
|
|
|
|
tf.saved_model.simple_save(sess, "./saved_model", inputs={"context": context}, outputs={"output": output})
|
|
|
|
def load_model():
|
|
fraction = 0.6
|
|
config = config = generate_gpu_config(fraction)
|
|
path_to_graph = "./saved"
|
|
|
|
# tf.saved_model.loader.load(
|
|
# session,
|
|
# [tf.saved_model.tag_constants.SERVING],
|
|
# path_to_graph)
|
|
|
|
# output = session.graph.get_tensor_by_name('output:0')
|
|
# context = session.graph.get_tensor_by_name('context:0')
|
|
model_path = 'gpt2/models/117M'
|
|
enc = encoder.get_encoder(model_path)
|
|
|
|
predict_fn = predictor.from_saved_model(path_to_graph, config=config)
|
|
context_tokens = [enc.encode("hello")]
|
|
predictions = predict_fn({"context": context_tokens})
|
|
output = enc.decode(predictions["output"][0])
|
|
print(output)
|
|
|
|
return (output, session)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
save_model()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|