From 185851ec959d99c7c92309781bf5212c5433dc05 Mon Sep 17 00:00:00 2001 From: Rose Hogenson Date: Tue, 15 Feb 2022 06:55:05 -0800 Subject: 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. --- ai_dungeon.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) (limited to 'ai_dungeon.py') 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:] -- cgit v1.3.1