update

console_play.py

@@ -23,12 +23,10 @@ def console_print(str, pycharm=False):
 def play_unconstrained():
     generator = CTRLGenerator()
     #generator = WebGenerator(CRED_FILE)
-    if len(sys.argv) is 1:
-        prompt = get_story_start("forest")
-    else:
-        prompt = sys.argv[1]
+    prompt = get_story_start("apocalypse")
+    context = get_context("apocalypse")
     story_manager = UnconstrainedStoryManager(generator)
-    story_manager.start_new_story(prompt)
+    story_manager.start_new_story(prompt, context=context)
 
     print("\n")
     console_print(str(story_manager.story))

generator/ctrl/ctrl_generator.py

@@ -20,7 +20,7 @@ def loss(labels, logits):
 
 class CTRLGenerator():
 
-    def __init__(self, control_code="Fantasy ", generate_num=100, temperature=0.5, topk=40, nucleus_prob=0):
+    def __init__(self, control_code="Writing Text:", generate_num=60, temperature=0.5, topk=40, nucleus_prob=0):
 
         self.generate_num=generate_num
         model_dir = "generator/ctrl/training_utils/seqlen256_v1.ckpt/"
@@ -158,7 +158,7 @@ class CTRLGenerator():
         if prompt[-1] != " ":
             prompt = prompt + " "
 
-        prompt = second_to_first_person(prompt)
+        #prompt = second_to_first_person(prompt)
 
         prompt = self.control_code + prompt
         # print("\n\nAFTER PROMPT_REPLACE")
@@ -172,7 +172,7 @@ class CTRLGenerator():
         first_letter_capitalized = result[0].isupper()
         result = result.replace("#", "")
         result = result.replace("*", "")
-        result = first_to_second_person(result)
+        #result = first_to_second_person(result)
         result = remove_profanity(result)
 
         if not first_letter_capitalized:
@@ -219,7 +219,7 @@ class CTRLGenerator():
         forbidden_tokens = ['<unk>', 'Sco@@', "&amp@@", "1]@@", "2]@@", "3]@@", "4]@@", "https://www.@@", "[@@", ":@@",
                             "Edit", "&@@", "2:","1:", ":", "Edit@@", "EDI@@", "EDIT@@", "edit", "TL@@", "tl@@", ";@@",
                             '**', "http://@@", "Redd@@", "UP@@", "mom", "Up@@", "Me:", "Update", "mom@@", "Part",
-                            "http://www.@@", "edit@@", "*@@"]
+                            "http://www.@@", "edit@@", "*@@", "\n", "Writing", "Text@@"]
 
         if num_new_lines > self.max_new_lines:
             forbidden_tokens.append("\n")

story/story_data.yaml

@@ -7,6 +7,8 @@ prompts:
 
   vague_police: "You are a police officer. You get to work and "
 
+  apocalypse: "You walk for two hours and take a break. You've left the town you were in and are now in a more rural area. There's a building to your right and you see"
+
 action_verbs:
   classic: ["You attack", "You tell", "You use", "You go"]
 
@@ -18,4 +20,7 @@ ctrl_verbs:
   movement: ["walk", "go", "run", "move"]
 
 rooms:
-  haunted_hospital: ["lobby", "hallway", "parking", "roof", "pharmacy"]
+  haunted_hospital: ["lobby", "hallway", "parking", "roof", "pharmacy"]
+
+contexts:
+  apocalypse: "Long ago the bombs fell and the world ended. You are one of the few who is still alive. You are trying to survive by scavenging among the ruins of what is left behind."

story/story_manager.py

@@ -5,8 +5,9 @@ import json
 
 class Story():
 
-    def __init__(self, story_start, seed=None, game_state=None):
+    def __init__(self, story_start, context ="", seed=None, game_state=None):
         self.story_start = story_start
+        self.context = context
 
         # list of actions. First action is the prompt length should always equal that of story blocks
         self.actions = []
@@ -33,18 +34,19 @@ class Story():
         self.choices = story_dict["choices"]
         self.possible_action_results = story_dict["possible_action_results"]
         self.game_state = story_dict["game_state"]
+        self.context = story_dict["context"]
 
     def add_to_story(self, action, story_block):
         self.actions.append(action)
         self.results.append(story_block)
 
     def latest_result(self):
-        if len(self.results) > 1:
-            return self.actions[-2] + self.results[-2] + self.actions[-1] + self.results[-1]
+        if len(self.results) >= 2:
+            return self.context + self.results[-1] + self.actions[-1] + self.results[-1]
         elif len(self.results) >= 1:
-            return self.story_start + self.actions[-1] + self.results[-1] 
+            return self.context + self.actions[-1] + self.results[-1]
         else:
-            return self.story_start
+            return self.context + self.story_start
 
     def __str__(self):
         story_list = [self.story_start]
@@ -63,7 +65,7 @@ class Story():
         story_dict["choices"] = self.choices
         story_dict["possible_action_results"] = self.possible_action_results
         story_dict["game_state"] = self.game_state
-
+        story_dict["context"] = self.context
 
         return json.dumps(story_dict)
 
@@ -72,10 +74,10 @@ class StoryManager():
     def __init__(self, generator):
         self.generator = generator
         
-    def start_new_story(self, story_prompt, game_state=None):
+    def start_new_story(self, story_prompt, context="", game_state=None):
         block = self.generator.generate(story_prompt)
         block = cut_trailing_sentence(block)
-        self.story = Story(story_prompt + block, game_state=game_state)
+        self.story = Story(story_prompt + block, context=context, game_state=game_state)
         return self.story
     
     def load_story(self, story, from_json=False):

story/utils.py

@@ -7,6 +7,12 @@ YAML_FILE = "story/story_data.yaml"
 from profanityfilter import ProfanityFilter
 pf = ProfanityFilter()
 
+def get_context(key):
+    with open(YAML_FILE, 'r') as stream:
+        data_loaded = yaml.safe_load(stream)
+
+    return data_loaded["context"][key]
+
 def get_story_start(key):
     with open(YAML_FILE, 'r') as stream:
         data_loaded = yaml.safe_load(stream)