Use an nltk tokenizer to trim sentence output.

It kind of feels bad to depend on nltk and hugging face, but we're doing
the best we can. Too bad that nltk isn't easier to install, too.

1 files changed, 8 insertions, 7 deletions

diff --git a/ai_dungeon.py b/ai_dungeon.py
index 63f6092..a73cc7e 100755
--- a/ai_dungeon.py
+++ b/ai_dungeon.py
@@ -1,23 +1,24 @@
 #!/usr/bin/env nix-shell
-#!nix-shell -i python3 -p "python3.withPackages (pkgs: with pkgs; [ transformers pytorch keras ])"
+#!nix-shell -i python3 -p "python3.withPackages (pkgs: with pkgs; [ keras nltk pytorch transformers ])"
 """AI Dungeon is a text-adventure style game powered by AI."""
 
 import argparse
 import re
 import sys
 
+from nltk import tokenize
 import transformers
 
 
-sentence_end = re.compile(r"[.?!(\"'][^.?!(\"']*$")
+sentence_end = re.compile(r"[.?!\"')]$")
 
 
 def trim_sentence(message: str) -> str:
     """Remove extra output after the last period."""
-    match = sentence_end.search(message)
-    if match:
-        return message[: match.start() + 1]
-    return message
+    sentences = tokenize.sent_tokenize(message)
+    if sentence_end.search(sentences[-1]) or len(sentences) == 1:
+        return message
+    return " ".join(sentences[:-1])
 
 
 def generate(model: transformers.TextGenerationPipeline, prompt: str, **kwargs) -> str:
@@ -129,7 +130,7 @@ def main() -> None:
             prev_response_lines = len(new_response.split("\n"))
             continue
         else:
-            if not msg.endswith("."):
+            if not sentence_end.search(msg):
                 msg += "."
             prompt += f"You {msg}\n"
         prompt = prompt[-10000:]