pass prompt to generator
This commit is contained in:
parent
943f351afe
commit
32313f9106
1 changed files with 3 additions and 3 deletions
|
@ -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()
|
||||
|
|
Loading…
Add table
Reference in a new issue