Simplify GPT adventure.

Or did I make it more complicated?

1 files changed, 104 insertions, 144 deletions

diff --git a/gpt_adventure.py b/gpt_adventure.py
index 65ff09c..ce77222 100755
--- a/gpt_adventure.py
+++ b/gpt_adventure.py
@@ -2,211 +2,171 @@
 #!nix-shell -i python3 -p "python3.withPackages (pkgs: with pkgs; [ keras pytorch transformers ])"
 """GPT adventure is a text-adventure style game powered by AI."""
 
+from typing import List, Optional
 import argparse
-import shutil
-import sys
 
 import transformers
 
 
-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,
-            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:
-    """Load a model by name."""
-    tokenizer = transformers.AutoTokenizer.from_pretrained(model)
-    return transformers.pipeline(
-        "text-generation",
-        tokenizer=tokenizer,
-        model=transformers.AutoModelForCausalLM.from_pretrained(model),
-        pad_token_id=tokenizer.eos_token_id,
-    )
+class Model:
+    def __init__(self, name: str):
+        self.name = name
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(name)
+        self.model = transformers.pipeline(
+            "text-generation",
+            tokenizer=self.tokenizer,
+            model=transformers.AutoModelForCausalLM.from_pretrained(name),
+            pad_token_id=self.tokenizer.eos_token_id,
+        )
 
+    def generate(self, prompt: str) -> str:
+        """Generate text from a text generator."""
+        out = self.model(
+                prompt,
+                do_sample=True,
+                temperature=0.9,
+                top_k=60,
+                top_p=0.9,
+                return_full_text=False,
+                max_new_tokens=50,
+        )[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"
 
-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"
+    def token_length(self, msg: str) -> int:
+        return len(self.tokenizer(msg)["input_ids"])
 
 
 class State:
-    def __init__(self, persona: Persona, history: str):
-        self.persona = persona
+    def __init__(self, model: Model, character: str, bio: str, history: List[str], filename: Optional[str]):
+        self.model = model
+        self.character = character
+        self.bio = bio
         self.history = history
-        self.prev_response_length = 0
+        self.filename = filename
+        self.prologue = f"{character}'s Persona: {bio}\n<START>\n"
 
     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)
+        self.history.append(f"{who}: {msg}")
 
     def rollback_history(self) -> None:
-        self.history = self.history[:-self.prev_response_length]
-        self.prev_response_length = 0
+        self.history = self.history[:-1]
 
     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}: "
+        return self.prologue + "\n".join(self.history) + f"\n{self.character}: "
+
+    def trim_history(self):
+        while self.history and self.model.token_length(self.prompt()) > 1024:
+            self.history = self.history[1:]
+
+    def print_context(self) -> None:
+        for elem in self.history[-10:]:
+            if elem.startswith("You: "):
+                print("> " + elem.removeprefix("You: "))
+            else:
+                print(elem.removeprefix(f"{self.character}: "))
+
+    def autosave(self) -> None:
+        if self.filename:
+            self.save(self.filename)
 
     def save(self, filename: str) -> None:
+        self.filename = filename
         with open(filename, "w") as f:
-            print(self.persona.name, file=f)
-            print(self.persona.bio, file=f)
-            print(self.history, file=f, end='')
+            print(self.model.name, file=f)
+            print(self.character, file=f)
+            print(self.bio, file=f)
+            for elem in self.history:
+                print(elem, file=f)
 
 
 def load(filename: str) -> State:
     with open(filename, "r") as f:
-        persona = f.readline().strip()
+        model = f.readline().strip()
+        character = f.readline().strip()
         bio = f.readline().strip()
-        history = f.read()
-    return State(Persona(persona, bio), history)
-
-
-def wrap(message: str) -> str:
-    """Wrap long lines to terminal width characters."""
-    width = shutil.get_terminal_size().columns
-    res = []
-    for line in message.split("\n"):
-        if len(line) < width:
-            res.append(line)
-            continue
-        split_line = []
-        for word in line.split(" "):
-            if not split_line or len(" ".join(split_line)) + len(word) < width:
-                split_line.append(word)
-                continue
-            res.append(" ".join(split_line))
-            split_line = [word]
-        res.append(" ".join(split_line))
-    return "\n".join(res)
-
-
-class Term:
-    def __init__(self):
-        self.rewind_point = 0
+        history = f.read().strip().split("\n")
+    return State(Model(model), character, bio, history, filename)
 
-    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 clear() -> None:
+    print(f"\033[2J", flush=True)
 
-    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 new_character(model: Model) -> State:
+    print("Creating a new character.")
+    print("You may have to heavily edit (/edit) the AI's first few responses")
+    print("in order for it to learn the appropriate tone.")
+    name = input("Enter your character's name: ")
+    bio = input("Enter a few sentences describing your character: ")
+    return State(model, name, bio, [], None)
 
 
 def main() -> None:
     """Run the main game loop."""
-    parser = argparse.ArgumentParser("AI dungeon clone")
+    parser = argparse.ArgumentParser("AI girlfriend")
     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)
-
     if args.save_file:
         state = load(args.save_file)
-        print("\n".join(state.history.rsplit("\n", 11)[-11:]).strip())
+        state.print_context()
     else:
-        state = pick_flavor()
+        state = new_character(Model(args.model))
 
-    term = Term()
     while True:
         try:
-            msg = term.input("> ")
+            msg = input("> ")
         except EOFError:
             break
-        if msg == "/retry":
+        if msg == "/help":
+            commands = (
+                    "/help",
+                    "/retry",
+                    "/edit",
+                    "/undo",
+                    "/save",
+            )
+            print("Commands:")
+            for cmd in commands:
+                print(f"\t{cmd}")
+        elif msg == "/retry":
             state.rollback_history()
-            term.rewind()
+            clear()
+            state.print_context()
         elif msg == "/edit":
             state.rollback_history()
-            term.rewind()
+            clear()
+            state.print_context()
 
-            new_response = term.input("")
-            state.add_history(state.persona.name, new_response)
+            new_response = input("")
+            state.add_history(state.character, new_response)
+            continue
+        elif msg == "/undo":
+            state.rollback_history()
+            state.rollback_history()
+            clear()
+            state.print_context()
             continue
         elif msg.startswith("/save"):
             if " " not in msg:
-                term.print("Usage: /save <filename>")
+                print("Usage: /save <filename>")
                 continue
             state.save(msg.split(" ", 1)[1])
             continue
         else:
             state.add_history("You", msg)
 
-        term.set_rewind_point()
+        state.trim_history()
 
-        response = generate(model, state.prompt(), max_new_tokens=50)
+        response = state.model.generate(state.prompt())
         if not response:
             response = "haha I'm not sure what you mean"
-        term.print(response)
-        state.add_history(state.persona.name, response)
+        print(response)
+        state.add_history(state.character, response)
 
 
 main()