diff --git a/console_play.py b/console_play.py index 729f646..b3794df 100644 --- a/console_play.py +++ b/console_play.py @@ -21,11 +21,12 @@ def console_print(str, pycharm=True): def play_unconstrained(): - #generator = CTRLGenerator() - generator = WebGenerator(CRED_FILE) - prompt = get_story_start("classic") + generator = CTRLGenerator() + #generator = WebGenerator(CRED_FILE) + prompt = get_story_start("haunted") story_manager = UnconstrainedStoryManager(generator, prompt) + print("\n") console_print(str(story_manager.story)) while (True): action = "" @@ -38,6 +39,7 @@ def play_unconstrained(): def play_constrained(): + print("\n") #generator = WebGenerator(CRED_FILE) generator = CTRLGenerator() story_start = "haunted" diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index 0f25415..aa7db7c 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -10,6 +10,8 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import embedding_ops import fastBPE from story.utils import * +import warnings +warnings.filterwarnings("ignore") pos_action_starts = ["You attack", "You tell", "You use", "You go"] @@ -20,7 +22,7 @@ def loss(labels, logits): class CTRLGenerator(): - def __init__(self, control_code="Horror Text: ", generate_num=64, temperature=0.3): + def __init__(self, control_code="Horror Text: ", generate_num=64, temperature=0.5): self.generate_num=generate_num model_dir = "generator/ctrl/model/seqlen256_v1.ckpt/" @@ -147,9 +149,9 @@ class CTRLGenerator(): prompt = self.control_code + prompt - # print("******************************") - # print(" DEBUG:: Prompt to generate by is \n", prompt) - # print("******************************") + print("******************************") + print(" DEBUG:: Prompt to generate by is \n", prompt) + print("******************************") prompt_length = len(prompt) @@ -162,6 +164,8 @@ class CTRLGenerator(): padded_text = text + [0] * (total_text_len - len(text)) tokens_generated = np.tile(padded_text, (1, 1)) result = "" + max_new_lines = 5 + num_new_lines = 0 for token in range(len(text) - 1, total_text_len - 1): if first_token: @@ -205,14 +209,18 @@ class CTRLGenerator(): prompt_logits[_token][generated_token] /= self.penalty # disallow some tokens - forbidden_tokens = ['', 'Sco@@', '\n'] + forbidden_tokens = ['', 'Sco@@'] + + if num_new_lines >= max_new_lines: + forbidden_tokens.append('\n') + for forbidden_token in forbidden_tokens: prompt_logits[_token][self.word2idx[forbidden_token]] = -1e8 # Make sure only a possible verb is chosen. if first_token: for word in get_possible_verbs(): - prompt_logits[_token][self.word2idx[word]] += 10 + prompt_logits[_token][self.word2idx[word]] += 5 # compute probabilities from logits prompt_probs = np.exp(prompt_logits[_token]) @@ -257,9 +265,12 @@ class CTRLGenerator(): # like which one was chosen, # what the top25 were, # here is your opportunity. - #print('chosen:', repr(self.idx2word[idx])) + print('chosen:', repr(self.idx2word[idx])) # print('top25 alternatives:', pruned_list[:25]) + if self.idx2word[idx] == "\n": + num_new_lines += 1 + # assign the token for generation tokens_generated[0][token + 1] = idx diff --git a/story/story_manager.py b/story/story_manager.py index 4d8d6a7..3997217 100644 --- a/story/story_manager.py +++ b/story/story_manager.py @@ -66,7 +66,6 @@ class UnconstrainedStoryManager(StoryManager): return result def generate_result(self, action): - print("Story context is ", self.story_context()) block = self.generator.generate(self.story_context() + action) block = cut_trailing_sentence(block) block = story_replace(block) diff --git a/story/utils.py b/story/utils.py index 34d5618..b946876 100644 --- a/story/utils.py +++ b/story/utils.py @@ -17,7 +17,7 @@ def get_action_verbs(key): return data_loaded["action_verbs"][key] - +# TODO add capital words to remove words def remove_profanity(text): remove_words = ["fuck", "Fuck", "shit", "rape", "bastard", "bitch"] for word in remove_words: @@ -87,7 +87,8 @@ first_to_second_mappings = [ ("we ","you "), ("We ","You "), (" mine"," yours"), - (" me", " you"), + (" me ", " you "), + (" me.", " you."), (" us ", " you "), (" us.", " you."), (" our", " your")