1 files changed, 22 insertions, 11 deletions
diff --git a/gpt_adventure.py b/gpt_adventure.py
index 68f8f58..c2137db 100755
--- a/gpt_adventure.py
+++ b/gpt_adventure.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:
@@ -68,6 +68,17 @@ def pick_flavor() -> str:
             "purifier for your settlement. You arrive at an\n"
             "abandoned settlement to the east of your home."
         ),
+        "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()):
@@ -114,18 +125,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