Add a new communism scenario.

1 files changed, 28 insertions, 17 deletions

diff --git a/ai_dungeon.py b/ai_dungeon.py
index b0cc64a..bd39133 100755
--- a/ai_dungeon.py
+++ b/ai_dungeon.py
@@ -14,14 +14,14 @@ import transformers
 
 def load_nltk():
     try:
-        tokenize.sent_tokenize('')
+        tokenize.sent_tokenize("")
     except LookupError:
         nltk.download("punkt")
 
 
 def complete_sentence(snippet: str) -> bool:
     sentences = tokenize.sent_tokenize(snippet)
-    return len(tokenize.sent_tokenize(sentences[-1] + ' extra')) == 2
+    return len(tokenize.sent_tokenize(sentences[-1] + " extra")) == 2
 
 
 def trim_sentence(message: str) -> str:
@@ -72,14 +72,25 @@ def pick_flavor() -> str:
             "Your name is Amelia Plenn. Today is your doctors appointment.\n"
             "You're kind of not looking forward to it because your doctor\n"
             "is kind of weird, but today when you get to the doctor's office\n"
-            "it's a different person than usual.\n\n"
-            "The new doctor says to you, \"Hello, my name is Doctor Cockman.\n"
-            "It's my pleasure to see you today.\"\n\n"
+            "it's a different person than usual.\n"
+            'The new doctor says to you, "Hello, my name is Doctor Cockman.\n'
+            "It's my pleasure to see you today.\"\n"
             "Doctor Cockman is very handsome, and you feel yourself blush as\n"
             "he leads you to the doctor's seat. You sit in the seat and he\n"
-            "runs his warm hands up and down your arms.\n\n"
-            "\"Are you feeling alright?\" Doctor Cockman asks."
-        )
+            "runs his warm hands up and down your arms.\n"
+            '"Are you feeling alright?" Doctor Cockman asks.'
+        ),
+        "political": (
+            "You are a political dissident and revolutionary named\n"
+            "Vladimir Lenin, hiding in exile in Finland while your homeland \n"
+            "of Russia is controlled by the Tsar. It is January 1917, and\n"
+            "The Great War has overtaken Europe. Millions of Russian peasants are\n"
+            "being sacrificed at the altar of international capital. You\n"
+            "are the leader of the political party Bolsheviks. You've been\n"
+            "writing letter after letter but the Bolsheviks in Russia don't\n"
+            "seem to want to publish your letters in Iskra. Maybe it's\n"
+            "time to return home."
+        ),
     }
     choices = []
     for i, (scenario, script) in enumerate(flavors.items()):
@@ -122,18 +133,18 @@ def main() -> None:
     """Run the main game loop."""
     load_nltk()
     parser = argparse.ArgumentParser("AI dungeon clone")
-    parser.add_argument(
-        "--model", default="gpt2", help="Model to use"
-    )
+    parser.add_argument("--model", default="gpt2", help="Model to use")
     args = parser.parse_args()
     model = load_model(args.model)
     script = pick_flavor()
-    prologue = wrap(generate(
-        model,
-        script,
-        max_new_tokens=20,
-        forced_eos_token_id=model.tokenizer.eos_token_id,
-    ))
+    prologue = wrap(
+        generate(
+            model,
+            script,
+            max_new_tokens=20,
+            forced_eos_token_id=model.tokenizer.eos_token_id,
+        )
+    )
     print(prologue)
     prompt = prologue
     prev_response_length = 0