From a906adb9ea2073ee9c23913585ff821558717ac8 Mon Sep 17 00:00:00 2001 From: Rose Hogenson Date: Tue, 15 Feb 2022 07:04:12 -0800 Subject: Auto-download nltk packages if they're missing. --- ai_dungeon.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) (limited to 'ai_dungeon.py') 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]) -- cgit v1.3.1