mirror of
https://github.com/wassname/Clover-Edition.git
synced 2026-09-25 13:10:22 +08:00
144 lines
5.0 KiB
Python
144 lines
5.0 KiB
Python
from story.utils import *
|
|
import warnings
|
|
warnings.filterwarnings("ignore")
|
|
import gpt_2_simple as gpt2
|
|
import os
|
|
import requests
|
|
import sys
|
|
import tensorflow as tf
|
|
import requests
|
|
from tqdm import tqdm
|
|
from gpt_2_simple.src import model, sample, encoder, memory_saving_gradients
|
|
from gpt_2_simple.src.load_dataset import load_dataset, Sampler
|
|
from gpt_2_simple.src.accumulate import AccumulatingOptimizer
|
|
import json
|
|
import numpy as np
|
|
|
|
tf.logging.set_verbosity(tf.logging.ERROR)
|
|
|
|
class SimpleGenerator:
|
|
|
|
def __init__(self, generate_num=65, temperature=0.9, top_k=40, top_p=0.0):
|
|
self.generate_num=generate_num
|
|
self.temp = temperature
|
|
self.top_k = top_k
|
|
self.top_p = top_p
|
|
|
|
self.model_name = "run1"
|
|
self.model_dir = "generator/simple/checkpoint"
|
|
self.checkpoint_path = os.path.join(self.model_dir, self.model_name)
|
|
if not os.path.isdir(os.path.join(self.model_dir, self.model_name)):
|
|
print(f"Downloading {self.model_name} model...")
|
|
|
|
subdir = os.path.join(self.model_dir, self.model_name)
|
|
if not os.path.exists(subdir):
|
|
os.makedirs(subdir)
|
|
subdir = subdir.replace('\\', '/') # needed for Windows
|
|
|
|
for filename in ['checkpoint', 'encoder.json', 'hparams.json', 'model.ckpt.data-00000-of-00001',
|
|
'model.ckpt.index', 'model.ckpt.meta', 'vocab.bpe']:
|
|
|
|
r = requests.get("https://storage.googleapis.com/gpt-2/" + subdir + "/" + filename, stream=True)
|
|
|
|
with open(os.path.join(subdir, filename), 'wb') as f:
|
|
file_size = int(r.headers["content-length"])
|
|
chunk_size = 1000
|
|
with tqdm(ncols=100, desc="Fetching " + filename, total=file_size, unit_scale=True) as pbar:
|
|
# 1k for chunk_size, since Ethernet packet size is around 1500 bytes
|
|
for chunk in r.iter_content(chunk_size=chunk_size):
|
|
f.write(chunk)
|
|
pbar.update(chunk_size)
|
|
|
|
self.sess = gpt2.start_tf_sess()
|
|
gpt2.load_gpt2(self.sess, model_dir=self.model_dir, model_name=self.model_name)
|
|
|
|
|
|
def prompt_replace(self, prompt):
|
|
# print("\n\nBEFORE PROMPT_REPLACE:")
|
|
# print(repr(prompt))
|
|
if prompt[-1] == " ":
|
|
prompt = prompt[:-1]
|
|
|
|
#prompt = second_to_first_person(prompt)
|
|
|
|
# print("\n\nAFTER PROMPT_REPLACE")
|
|
# print(repr(prompt))
|
|
return prompt
|
|
|
|
def result_replace(self, result):
|
|
# print("\n\nBEFORE RESULT_REPLACE:")
|
|
# print(repr(result))
|
|
|
|
result = cut_trailing_sentence(result)
|
|
first_letter_capitalized = result[0].isupper()
|
|
result = result.replace('."', '".')
|
|
result = result.replace("#", "")
|
|
result = result.replace("*", "")
|
|
#result = first_to_second_person(result)
|
|
result = remove_profanity(result)
|
|
|
|
if not first_letter_capitalized:
|
|
result = result[0].lower() + result[1:]
|
|
|
|
#
|
|
# print("\n\nAFTER RESULT_REPLACE:")
|
|
# print(repr(result))
|
|
|
|
return result
|
|
|
|
def generate(self, prompt, options=None, seed=1):
|
|
|
|
debug=True
|
|
prefix = self.prompt_replace(prompt)
|
|
|
|
if debug:
|
|
print("******DEBUG******")
|
|
print("Prompt is: ", repr(prefix))
|
|
|
|
|
|
enc = encoder.get_encoder(self.checkpoint_path)
|
|
hparams = model.default_hparams()
|
|
batch_size = 1
|
|
nsamples = 1
|
|
with open(os.path.join(self.checkpoint_path, 'hparams.json')) as f:
|
|
hparams.override_from_dict(json.load(f))
|
|
|
|
context = tf.compat.v1.placeholder(tf.int32, [batch_size, None])
|
|
context_tokens = enc.encode(prefix)
|
|
|
|
#np.random.seed(seed)
|
|
#tf.compat.v1.set_random_seed(seed)
|
|
|
|
output = sample.sample_sequence(
|
|
hparams=hparams,
|
|
length=min(self.generate_num, 1023 - (len(context_tokens) if prefix else 0)),
|
|
start_token=enc.encoder['<|endoftext|>'] if not prefix else None,
|
|
context=context if prefix else None,
|
|
batch_size=batch_size,
|
|
temperature=self.temp, top_k=self.top_k, top_p=self.top_p
|
|
)[:, 1:]
|
|
|
|
generated = 0
|
|
gen_texts = []
|
|
while generated < nsamples:
|
|
if not prefix:
|
|
out = self.sess.run(output)
|
|
else:
|
|
out = self.sess.run(output, feed_dict={
|
|
context: batch_size * [context_tokens]
|
|
})
|
|
for i in range(batch_size):
|
|
generated += 1
|
|
gen_text = enc.decode(out[i])
|
|
if prefix:
|
|
gen_text = enc.decode(context_tokens[:1]) + gen_text
|
|
gen_text = gen_text.lstrip('\n')
|
|
gen_texts.append(gen_text)
|
|
|
|
print("Generated result is: ", repr(gen_texts[0]))
|
|
print("******END DEBUG******")
|
|
result = gen_texts[0][len(prefix):]
|
|
|
|
result = self.result_replace(result)
|
|
return result
|