From a4d6412fdaad7fae011d7103d4e40b01091b112d Mon Sep 17 00:00:00 2001
From: Rose Hogenson <rhogenson@posteo.net>
Date: Tue, 15 Feb 2022 09:02:56 -0800
Subject: Wrap the prologue.

Sometimes the AI generates really long lines.
---
 ai_dungeon.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'ai_dungeon.py')

diff --git a/ai_dungeon.py b/ai_dungeon.py
index 2b976f6..baac638 100755
--- a/ai_dungeon.py
+++ b/ai_dungeon.py
@@ -110,12 +110,12 @@ def main() -> None:
     args = parser.parse_args()
     model = load_model(args.model)
     script = pick_flavor()
-    prologue = generate(
+    prologue = wrap(generate(
         model,
         script,
         max_new_tokens=20,
         forced_eos_token_id=model.tokenizer.eos_token_id,
-    )
+    ))
     print(prologue)
     prompt = prologue
     prev_response_length = 0
-- 
cgit v1.3.1