From 9d2837981c40df7a5b179e3f217e68df3f210035 Mon Sep 17 00:00:00 2001 From: Nick Date: Thu, 19 Sep 2019 00:02:52 -0600 Subject: [PATCH] update --- generator/ctrl/ctrl_generator.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index b976510..0e58604 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -9,6 +9,7 @@ from tensorflow.python import debug as tf_debug from tensorflow.python.ops import math_ops from tensorflow.python.ops import embedding_ops import fastBPE +from story.utils import * pos_action_starts = ["You attack", "You tell", "You use", "You go"] @@ -137,7 +138,10 @@ class CTRLGenerator(): def generate(self, prompt): + prompt = second_to_first_person(prompt) + prompt = self.control_code + prompt + prompt_length = len(prompt) # tokenize provided prompt split_prompt = self.bpe.apply([prompt])[0].split() @@ -248,7 +252,7 @@ class CTRLGenerator(): tokens_generated_so_far = re.sub('(@@ )', '', string=tokens_generated_so_far) tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far) - result = tokens_generated_so_far + result = tokens_generated_so_far[prompt_length:] except: print("Error in generation")