diff options
| -rwxr-xr-x | ai_dungeon.py | 7 |
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]) |