1 files changed, 57 insertions, 60 deletions
diff --git a/ai_dungeon.py b/ai_dungeon.py
index f1aabf4..ac5fe82 100755
--- a/ai_dungeon.py
+++ b/ai_dungeon.py
@@ -9,138 +9,135 @@ import transformers
 
 def trim_sentence(tokenizer: transformers.PreTrainedTokenizer, message: str) -> str:
     """Remove extra output after the last period."""
-    periodt = tokenizer.encode('.')[0]
+    periodt = tokenizer.encode(".")[0]
     tokens = tokenizer.encode(message)
     for i in range(len(tokens) - 1, -1, -1):
         if tokens[i] == periodt:
-            return tokenizer.decode(tokens[:i+1])
+            return tokenizer.decode(tokens[: i + 1])
     return message
 
 
 def generate(model: transformers.TextGenerationPipeline, prompt: str, **kwargs) -> str:
     """Generate text from a text generator."""
-    out = model(
-            prompt,
-            do_sample=True,
-            temperature=0.8,
-            top_k=60,
-            top_p=0.9,
-            **kwargs)
-    return trim_sentence(model.tokenizer, out[0]['generated_text'])
+    out = model(prompt, do_sample=True, temperature=0.8, top_k=60, top_p=0.9, **kwargs)
+    return trim_sentence(model.tokenizer, out[0]["generated_text"])
 
 
 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)
+        "text-generation",
+        tokenizer=tokenizer,
+        model=transformers.AutoModelForCausalLM.from_pretrained(model),
+        pad_token_id=tokenizer.eos_token_id,
+    )
 
 
 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."
+        "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."
+        ),
     }
     choices = []
     for i, (scenario, script) in enumerate(flavors.items()):
-        print(f'{i}.\t{scenario}')
+        print(f"{i}.\t{scenario}")
         choices.append(script)
     while True:
-        choice = input('Choose a scenario: ')
+        choice = input("Choose a scenario: ")
         try:
             choice_int = int(choice)
         except ValueError:
-            print('Input must be an integer.')
+            print("Input must be an integer.")
             continue
         try:
             return choices[choice_int]
         except IndexError:
-            print('Index out of bounds.')
+            print("Index out of bounds.")
             continue
 
 
 def wrap(message: str) -> str:
     """Wrap long lines to 72 characters."""
     res = []
-    for line in message.split('\n'):
+    for line in message.split("\n"):
         if len(line) < 72:
             res.append(line)
             continue
         split_line = []
-        for word in line.split(' '):
-            if not split_line or len(' '.join(split_line)) + len(word) < 72:
+        for word in line.split(" "):
+            if not split_line or len(" ".join(split_line)) + len(word) < 72:
                 split_line.append(word)
                 continue
-            res.append(' '.join(split_line))
+            res.append(" ".join(split_line))
             split_line = [word]
-        res.append(' '.join(split_line))
-    return '\n'.join(res)
+        res.append(" ".join(split_line))
+    return "\n".join(res)
 
 
 def main() -> None:
     """Run the main game loop."""
-    parser = argparse.ArgumentParser('AI dungeon clone')
+    parser = argparse.ArgumentParser("AI dungeon clone")
     parser.add_argument(
-            '--model',
-            default='EleutherAI/gpt-neo-125M',
-            help='Model to use')
+        "--model", default="EleutherAI/gpt-neo-125M", help="Model to use"
+    )
     args = parser.parse_args()
     model = load_model(args.model)
     script = pick_flavor()
     prologue = generate(
-            model,
-            script,
-            max_new_tokens=20,
-            forced_eos_token_id=model.tokenizer.eos_token_id)
+        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
     while True:
         try:
-            msg = input('> You ')
+            msg = input("> You ")
         except EOFError:
             break
-        if msg == '/retry':
+        if msg == "/retry":
             prompt = prompt[:-prev_response_length]
-            print(f'\033[{prev_response_lines}A\033[J', end='')
-        elif msg == '/edit':
+            print(f"\033[{prev_response_lines}A\033[J", end="")
+        elif msg == "/edit":
             prompt = prompt[:-prev_response_length]
-            print(f'\033[{prev_response_lines}A\033[J\r', end='')
+            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'))
+            prev_response_lines = len(new_response.split("\n"))
             continue
         else:
-            if not msg.endswith('.'):
-                msg += '.'
-            prompt += f'You {msg}\n'
+            if not msg.endswith("."):
+                msg += "."
+            prompt += f"You {msg}\n"
         prompt = prompt[-10000:]
-        response = wrap(generate(
-                model,
-                prompt,
-                return_full_text=False,
-                max_new_tokens=50).strip())
+        response = wrap(
+            generate(model, prompt, return_full_text=False, max_new_tokens=50).strip()
+        )
         print(response)
-        response += '\n'
+        response += "\n"
         prompt += response
         prev_response_length = len(response)
-        prev_response_lines = len(response.split('\n'))
+        prev_response_lines = len(response.split("\n"))
 
 
 main()