update

generator/ctrl/ctrl_generator.py

@@ -297,7 +297,7 @@ class CTRLGenerator():
         num_new_lines = 0
         for token in range(len(text) - 1, total_text_len - 1):
             idx = self.generate_next_token(token, tokens_generated, options, num_new_lines, token_num, first_token=first_token, forbid_newline=False)
-            if self.idx2word[idx] == '\n' and token_num < 10:
+            if self.idx2word[idx] == '\n' and token_num > 7:
                 return self.result_replace(result)
             elif self.idx2word[idx] == '\n':
                 idx = self.generate_next_token(token, tokens_generated, options, num_new_lines, token_num,

story/story_manager.py

@@ -209,6 +209,7 @@ class CTRLStoryManager(ConstrainedStoryManager):
     def get_action_results_generate(self):
         results = []
         options = {"word_blacklist": {0:[]}}
+        options["word_whitelist"] = {0: get_allowed_ctrl_verbs()}
         for phrase in self.action_phrases:
             result = self.generate_action_result(self.story_context(), phrase, options=options)
             action_verb = result[0].split()[1]

story/utils.py

@@ -13,6 +13,12 @@ def get_context(key):
 
     return data_loaded["contexts"][key]
 
+def get_allowed_ctrl_verbs():
+    with open(YAML_FILE, 'r') as stream:
+        data_loaded = yaml.safe_load(stream)
+
+    return data_loaded["ctrl_verbs"]["movement"] + data_loaded["ctrl_verbs"]["non_movement"]
+
 def get_story_start(key):
     with open(YAML_FILE, 'r') as stream:
         data_loaded = yaml.safe_load(stream)