Redo GPT adventure to be a chat bot.

I found a cool transformer that works really well for chat bots. So now
instead of a GPT-adventure style RPG, you're chatting with an
AI girlfriend.

1 files changed, 150 insertions, 114 deletions

diff --git a/gpt_adventure.py b/gpt_adventure.py
index 6fda9fe..9f93d5f 100755
--- a/gpt_adventure.py
+++ b/gpt_adventure.py
@@ -1,42 +1,29 @@
 #!/usr/bin/env nix-shell
-#!nix-shell -i python3 -p "python3.withPackages (pkgs: with pkgs; [ keras nltk pytorch transformers ])"
+#!nix-shell -i python3 -p "python3.withPackages (pkgs: with pkgs; [ keras pytorch transformers ])"
 """GPT adventure is a text-adventure style game powered by AI."""
 
 import argparse
 import shutil
 import sys
 
-import nltk
-from nltk import tokenize
 import transformers
 
 
-def load_nltk():
-    """Download required nltk modules."""
-    try:
-        tokenize.sent_tokenize("")
-    except LookupError:
-        nltk.download("punkt")
-
-
-def complete_sentence(snippet: str) -> bool:
-    """Return whether the snippet ends in a complete sentence."""
-    sentences = tokenize.sent_tokenize(snippet)
-    return len(tokenize.sent_tokenize(sentences[-1] + " extra")) == 2
-
-
-def trim_sentence(message: str) -> str:
-    """Remove extra output after the last period."""
-    sentences = tokenize.sent_tokenize(message)
-    if complete_sentence(sentences[-1]) or len(sentences) == 1:
-        return message
-    return " ".join(sentences[:-1])
-
-
 def generate(model: transformers.TextGenerationPipeline, prompt: str, **kwargs) -> str:
     """Generate text from a text generator."""
-    out = model(prompt, do_sample=True, temperature=0.9, top_k=60, top_p=0.9, **kwargs)
-    return trim_sentence(out[0]["generated_text"])
+    out = model(
+            prompt,
+            do_sample=True,
+            temperature=0.9,
+            top_k=60,
+            top_p=0.9,
+            return_full_text=False,
+            **kwargs,
+    )[0]["generated_text"]
+    first_line = out.split("\n", 1)[0].strip()
+    if first_line:
+        return first_line
+    return "haha i'm not sure what you mean"
 
 
 def load_model(model: str) -> transformers.TextGenerationPipeline:
@@ -50,57 +37,47 @@ def load_model(model: str) -> transformers.TextGenerationPipeline:
     )
 
 
-def pick_flavor() -> str:
-    """Query the user for what scenario they want to play."""
-    flavors = {
-        "fantasy": (
-            "You are a wizard named Megumin from the kingdom\n"
-            "of Larion. You have in your inventory a wizard's staff\n"
-            "and a spellbook. You are arriving after a day's travel\n"
-            "at an enchanted tower where there's rumors of gold.\n"
-            "You walk up to the entrance of the tower."
-        ),
-        "post apocalyptic": (
-            "You are a machinist named Azariel, living\n"
-            "in the city of New New York. It's been almost 10 years\n"
-            "since the bombs fell, but you still remember it as if\n"
-            "it were yesterday. You push these thoughts out of your\n"
-            "mind and focus on the task at hand: finding a water\n"
-            "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()):
-        print(f"{i}.\t{scenario}")
-        choices.append(script)
-    print(f"{len(choices)}.\tcustom")
-    while True:
-        choice = input("Choose a scenario: ")
-        try:
-            choice_int = int(choice)
-        except ValueError:
-            print("Input must be an integer.")
-            continue
-        if choice_int == len(choices):
-            print("Enter a custom prompt. Press control-D when you're done.")
-            return sys.stdin.read().strip()
-        try:
-            return choices[choice_int]
-        except IndexError:
-            print("Index out of bounds.")
-            continue
+class Persona:
+    def __init__(self, name: str, bio: str):
+        self.name = name
+        self.bio = bio
+        self.prologue = f"{name}'s Persona: {bio}\n<START>\n"
+
+
+class State:
+    def __init__(self, persona: Persona, history: str):
+        self.persona = persona
+        self.history = history
+        self.prev_response_length = 0
+
+    def add_history(self, who: str, msg: str) -> None:
+        fmt_msg = f"{who}: {msg}\n"
+        self.history += fmt_msg
+        self.prev_response_length = len(fmt_msg)
+
+    def rollback_history(self) -> None:
+        self.history = self.history[:-self.prev_response_length]
+        self.prev_response_length = 0
+
+    def prompt(self) -> str:
+        max_prompt_size = 10000
+        if len(self.history) > max_prompt_size:
+            self.history = self.history[-(max_prompt_size+1):].split("\n", 1)[1]
+        return f"{self.persona.prologue}{self.history}{self.persona.name}: "
+
+    def save(self, filename: str) -> None:
+        with open(filename, "w") as f:
+            print(self.persona.name, file=f)
+            print(self.persona.bio, file=f)
+            print(self.history, file=f, end='')
+
+
+def load(filename: str) -> State:
+    with open(filename, "r") as f:
+        persona = f.readline().strip()
+        bio = f.readline().strip()
+        history = f.read()
+    return State(Persona(persona, bio), history)
 
 
 def wrap(message: str) -> str:
@@ -122,55 +99,114 @@ def wrap(message: str) -> str:
     return "\n".join(res)
 
 
+class Term:
+    def __init__(self):
+        self.rewind_point = 0
+
+    def input(self, prompt: str) -> str:
+        self.rewind_point += 1
+        return input(prompt)
+
+    def print(self, msg: str) -> None:
+        wrapped_msg = wrap(msg)
+        print(wrapped_msg)
+        self.rewind_point += wrapped_msg.count("\n") + 1
+
+    def set_rewind_point(self) -> None:
+        self.rewind_point = 0
+
+    def rewind(self) -> None:
+        print(f"\033[{self.rewind_point}A\033[J\r", end="")
+        self.set_rewind_point()
+
+
+flavors = (
+    State(Persona("Rin",
+        "Tohsaka Rin is a feisty and independent mage who can come across as "
+        "rude and unlikeable at first. She can be sweet and caring, but it "
+        "takes a lot to break down her guard."),
+
+        "Rin: What are you looking at... idiot?\n"
+        "You: Umm nothing...\n"
+        "Rin: That's right, you're nothing. You are less than a piece of trash.\n"
+        "You: Rin, do you want to go demon-hunting some time?\n"
+        "Rin: Well, maybe. But not because I like you or anything.\n"
+        "You: Right, we'll just go as friends.\n"
+        "Rin: Or maybe acquaintances...\n"
+    ),
+)
+
+
+def pick_flavor() -> State:
+    """Query the user for what scenario they want to play."""
+    for i, state in enumerate(flavors):
+        print(f"{i}.\t{state.persona.name}")
+    print(f"{len(flavors)}.\tcustom")
+    while True:
+        choice = input("Choose a persona: ")
+        try:
+            choice_int = int(choice)
+        except ValueError:
+            print("Input must be an integer.")
+            continue
+        if choice_int == len(flavors):
+            name = input("Enter your character's name: ")
+            bio = input("Enter your character's persona: ")
+            return State(Persona(name, bio), [])
+        try:
+            return flavors[choice_int]
+        except IndexError:
+            print("Index out of bounds.")
+            continue
+
+
 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="PygmalionAI/pygmalion-6b", help="Model to use")
+    parser.add_argument("save_file", nargs="?", default="")
     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,
-        )
-    )
-    print(prologue)
-    prompt = prologue
-    prev_response_length = 0
-    prev_response_lines = 0
+
+    if args.save_file:
+        state = load(args.save_file)
+        print("\n".join(state.history.rsplit("\n", 11)[-11:]).strip())
+    else:
+        state = pick_flavor()
+
+    term = Term()
     while True:
         try:
-            msg = input("> You ")
+            msg = term.input("> ")
         except EOFError:
             break
         if msg == "/retry":
-            prompt = prompt[:-prev_response_length]
-            print(f"\033[{prev_response_lines}A\033[J", end="")
+            state.rollback_history()
+            term.rewind()
         elif msg == "/edit":
-            prompt = prompt[:-prev_response_length]
-            print(f"\033[{prev_response_lines}A\033[J\r", end="")
-            new_response = sys.stdin.read()
-            prompt += new_response
-            prev_response_length = len(new_response)
-            prev_response_lines = len(new_response.split("\n"))
+            state.rollback_history()
+            term.rewind()
+
+            new_response = term.input("")
+            state.add_history(state.persona.name, new_response)
+            continue
+        elif msg.startswith("/save"):
+            if " " not in msg:
+                term.print("Usage: /save <filename>")
+                continue
+            state.save(msg.split(" ", 1)[1])
             continue
         else:
-            if not complete_sentence(msg):
-                msg += "."
-            prompt += f"You {msg}\n"
-        prompt = prompt[-10000:]
-        response = wrap(
-            generate(model, prompt, return_full_text=False, max_new_tokens=50).strip()
-        )
-        print(response)
-        response += "\n"
-        prompt += response
-        prev_response_length = len(response)
-        prev_response_lines = len(response.split("\n"))
+            state.add_history("You", msg)
+
+        term.set_rewind_point()
+
+        response = generate(model, state.prompt(), max_new_tokens=50)
+        if not response:
+            response = "haha I'm not sure what you mean"
+        term.print(response)
+        state.add_history(state.persona.name, response)
 
 
 main()