update

generator/ctrl/ctrl_generator.py

@@ -207,7 +207,7 @@ class CTRLGenerator():
             # Make sure only a possible verb is chosen.
             if first_token:
                 for word in get_possible_verbs():
-                    prompt_logits[_token][self.word2idx[word]] *= 1000
+                    prompt_logits[_token][self.word2idx[word]] += max(prompt_logits)
                 first_token = False
 
             # compute probabilities from logits

story/utils.py

@@ -102,6 +102,7 @@ second_to_first_mappings = [
     ("You are ", "I am "),
     ("You ", "I "),
     ("you ", "I "),
+    (" you.", " me."),
 ]
 
 def capitalize_helper(string):