From da015e21497a4dd7f11973bd832e7ed1cfe3398a Mon Sep 17 00:00:00 2001
From: Nick <nickwalton00@gmail.com>
Date: Thu, 19 Sep 2019 22:44:08 -0600
Subject: [PATCH] update

---
 console_play.py                  | 3 ++-
 generator/ctrl/ctrl_generator.py | 2 ++
 story/utils.py                   | 5 +++++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/console_play.py b/console_play.py
index e7bab80..b5c5d7b 100644
--- a/console_play.py
+++ b/console_play.py
@@ -29,8 +29,9 @@ def play_unconstrained():
 
     console_print(str(story_manager.story))
     while (True):
+        print("\n")
         action = input("> ")
-        action = "You " + action
+        action = " You " + action + "."
         result = story_manager.act(action)
         console_print(action + result)
 
diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py
index fec60f1..8b585ad 100644
--- a/generator/ctrl/ctrl_generator.py
+++ b/generator/ctrl/ctrl_generator.py
@@ -256,6 +256,8 @@ class CTRLGenerator():
             # os.system("clear")
 
             tokens_generated_so_far = ' '.join([self.idx2word[c] for c in tokens_generated[0].squeeze()[:token + 2]])
+            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[prompt_length:]
 
diff --git a/story/utils.py b/story/utils.py
index 750bd5c..0cab888 100644
--- a/story/utils.py
+++ b/story/utils.py
@@ -61,16 +61,21 @@ def replace_outside_quotes(text, current_word, repl_word):
 
 
 first_to_second_mappings = [
+    ("I'm", "you're"),
     ("I am", "you are"),
     ("I ", "you "),
     ("I've", "You've"),
     ("my", "your"),
+    ("My", "Your"),
     ("we ","you "),
     ("We ","You "),
     (" mine"," yours"),
+    ("me", "you"),
 ]
 
 second_to_first_mappings = [
+    ("You're", "I'm"),
+    ("you're", "I'm"),
     ("Your", "My"),
     ("your", "my"),
     ("you are", "I am"),