From 32313f910698b5bad7b264cc4753341c40deaa20 Mon Sep 17 00:00:00 2001
From: array-in-a-matrix <nova.neutrino@protonmail.com>
Date: Wed, 17 Aug 2022 16:24:53 -0400
Subject: [PATCH] pass prompt to generator

---
 textgen.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/textgen.py b/textgen.py
index 437bd5e..f26bfbb 100644
--- a/textgen.py
+++ b/textgen.py
@@ -13,9 +13,9 @@ file_name = json_object['file']
 # ? generate message using trained model
 
 
-def generate_message():
+def generate_message(prompt):
     ai = aitextgen(model_folder="trained_model",
-                   tokenizer_file="aitextgen.tokenizer.json")
+                   tokenizer_file="aitextgen.tokenizer.json", prompt=prompt)
     ai.generate()
 
 # ? train model using text file
@@ -35,7 +35,7 @@ def train_ai():
 match sys.argv[1]:
     case "generate":
         # ? send message to parent JS process
-        print(generate_message())
+        print(generate_message(sys.argv[2]))
         sys.stdout.flush()
     case "train":
         train_ai()