Auto-download nltk packages if they're missing.

1 files changed, 6 insertions, 1 deletions

diff --git a/ai_dungeon.py b/ai_dungeon.py
index a73cc7e..58ad206 100755
--- a/ai_dungeon.py
+++ b/ai_dungeon.py
@@ -6,6 +6,7 @@ import argparse
 import re
 import sys
 
+import nltk
 from nltk import tokenize
 import transformers
 
@@ -15,7 +16,11 @@ sentence_end = re.compile(r"[.?!\"')]$")
 
 def trim_sentence(message: str) -> str:
     """Remove extra output after the last period."""
-    sentences = tokenize.sent_tokenize(message)
+    try:
+        sentences = tokenize.sent_tokenize(message)
+    except LookupError:
+        nltk.download("punkt")
+        sentences = tokenize.sent_tokenize(message)
     if sentence_end.search(sentences[-1]) or len(sentences) == 1:
         return message
     return " ".join(sentences[:-1])