diff options
| -rwxr-xr-x | gpt_adventure.py | 33 |
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 |