From 60cd7954af55e91c97e90a1d8ae3e051c93dded5 Mon Sep 17 00:00:00 2001 From: Rose Hogenson Date: Tue, 15 Feb 2022 10:12:40 -0800 Subject: Use a better sentence detection strategy. It's actually pretty easy to see whether a snippet is a complete sentence, just append an extra word and see if it tokenizes into two sentences or just one. This should be a lot more reliable. --- ai_dungeon.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/ai_dungeon.py b/ai_dungeon.py index baac638..ad18c48 100755 --- a/ai_dungeon.py +++ b/ai_dungeon.py @@ -12,17 +12,22 @@ from nltk import tokenize import transformers -sentence_end = re.compile(r"[.?!\"')]$") +def load_nltk(): + try: + tokenize.sent_tokenize('') + except LookupError: + nltk.download("punkt") + + +def complete_sentence(snippet: str) -> bool: + sentences = tokenize.sent_tokenize(snippet) + return len(tokenize.sent_tokenize(sentences[-1] + ' extra')) == 2 def trim_sentence(message: str) -> str: """Remove extra output after the last period.""" - 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: + sentences = tokenize.sent_tokenize(message) + if complete_sentence(sentences[-1]) or len(sentences) == 1: return message return " ".join(sentences[:-1]) @@ -103,6 +108,7 @@ def wrap(message: str) -> str: def main() -> None: """Run the main game loop.""" + load_nltk() parser = argparse.ArgumentParser("AI dungeon clone") parser.add_argument( "--model", default="gpt2", help="Model to use" @@ -137,7 +143,7 @@ def main() -> None: prev_response_lines = len(new_response.split("\n")) continue else: - if not sentence_end.search(msg): + if not complete_sentence(msg): msg += "." prompt += f"You {msg}\n" prompt = prompt[-10000:] -- cgit v1.3.1