diff options
| -rwxr-xr-x | ai_dungeon.py | 15 |
1 files changed, 8 insertions, 7 deletions
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:] |