Commit 70e0c83a authored by Ilya Prokhorov's avatar Ilya Prokhorov

Dataset updated, command (batch) mode added

parent 06040baa
This diff is collapsed.
from textgenrnn import textgenrnn
from os import path
print("Textgenrnn quotes text generator by demensdeum 2018 (demensdeum@gmail.com)")
state = input("train/generate? ")
import sys
weights_filename = "textgenrnn_weights.hdf5"
vocab_filename = "textgenrnn_vocab.json"
config_filename = "textgenrnn_config.json"
default_dataset_filename = "dataset_quotes_ru.txt"
if state == "train":
reset_model = False
print("Textgenrnn quotes text generator by demensdeum 2018 (demensdeum@gmail.com)")
train_state = "reset"
if path.exists(weights_filename):
train_state = input("reset/resume? ")
command_mode = False
if len(sys.argv) == 2:
print("Silent mode enabled")
command_mode = True
command_mode_state = sys.argv[1]
command_mode_state_array = command_mode_state.split(",")
if len(command_mode_state_array) != 2:
print("Incorrect silent mode, must be - mode,dataset_filename - for example: python main.py generate,dataset_quotes_ru.txt")
print("Modes: train_reset/train_resume/generate")
if command_mode == False:
state = input("train_reset/train_resume/generate? ")
else:
state = command_mode_state_array[0]
if state != "train_reset" and state != "train_resume" and state != "generate":
print("Write train_reset or train_resume or generate... Exit")
exit(1)
if state == "train_reset":
textgen = textgenrnn()
if train_state == "reset":
reset_model = True
textgen = textgenrnn()
elif train_state == "resume":
reset_model = False
textgen = textgenrnn(weights_filename)
else:
print("Write reset or result... Exit")
exit(2)
elif state == "train_resume":
if path.exists(weights_filename) == False:
print("There is no weights to generate, train first... Exit")
exit(2)
textgen = textgenrnn(weights_path = weights_filename, vocab_path = vocab_filename, config_path = config_filename)
if command_mode == False and (state == "train_reset" or state == "train_resume"):
dataset_file = input("dataset filename? (%s) " % default_dataset_filename)
if len(dataset_file) < 1:
dataset_file = default_dataset_filename
while True:
print("Endless train mode, every 4 epochs will be saved. CTRL+C to exit")
textgen.train_from_file(dataset_file, num_epochs=4, new_model = reset_model)
dataset_file = default_dataset_filename
elif command_mode == True:
dataset_file = command_mode_state_array[1]
if state == "train_reset" or state == "train_resume":
reset_model = state == "train_reset"
while True:
print("Endless train mode, every 4 epochs will be saved. CTRL+C to exit")
try:
textgen.train_from_file(dataset_file, num_epochs=4, new_model = reset_model)
except KeyboardInterrupt:
print("\nKilled")
exit(0)
except:
print("Crashed... Recovery")
elif state == "generate":
if path.exists(weights_filename) == False:
print("There is no weights to generate, train first... Exit")
exit(3)
if path.exists(weights_filename) == False:
print("There is no weights to generate, train first... Exit")
exit(3)
textgen = textgenrnn(weights_path = weights_filename, vocab_path = vocab_filename,
config_path = config_filename)
textgen.generate()
else:
print("Write train or generate... Exit")
exit(1)
\ No newline at end of file
textgen = textgenrnn(weights_path = weights_filename, vocab_path = vocab_filename, config_path = config_filename)
textgen.generate()
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment