array-in-a-matrix/text-gen-bot
1
0
Fork 0
Code Issues 2 Pull requests Projects Releases Wiki Activity

check arg passed

Browse source
This commit is contained in:
array-in-a-matrix 2022-08-16 03:46:05 -04:00
parent 8c6ad6d5a3
commit 1b4b769b7a
1 changed files with 21 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
textgen.py
View file

@ -2,7 +2,8 @@ from aitextgen.TokenDataset import TokenDataset
from aitextgen.tokenizers import train_tokenizer from aitextgen.tokenizers import train_tokenizer
from aitextgen.utils import GPT2ConfigCPU from aitextgen.utils import GPT2ConfigCPU
from aitextgen import aitextgen from aitextgen import aitextgen
import json, sys import json
import sys
with open('config.json', 'r') as file: with open('config.json', 'r') as file:
json_object = json.load(file) json_object = json.load(file)
@ -10,21 +11,33 @@ with open('config.json', 'r') as file:
file_name = json_object['file'] file_name = json_object['file']
# ? generate message using trained model # ? generate message using trained model
def generate_message(): def generate_message():
ai = aitextgen(model_folder="trained_model", ai = aitextgen(model_folder="trained_model",
tokenizer_file="aitextgen.tokenizer.json") tokenizer_file="aitextgen.tokenizer.json")
ai.generate() ai.generate()
# ? train model using text file # ? train model using text file
def train_ai(): def train_ai():
train_tokenizer(file_name) train_tokenizer(file_name)
tokenizer_file = "aitextgen.tokenizer.json" tokenizer_file = "aitextgen.tokenizer.json"
config = GPT2ConfigCPU() config = GPT2ConfigCPU()
ai = aitextgen(tokenizer_file=tokenizer_file, config=config) ai = aitextgen(tokenizer_file=tokenizer_file, config=config)
data = TokenDataset(file_name, tokenizer_file=tokenizer_file, block_size=64) data = TokenDataset(
ai.train(data, batch_size=8, num_steps=50000, generate_every=5000, save_every=5000) file_name, tokenizer_file=tokenizer_file, block_size=64)
print("AI has been trained!") ai.train(data, batch_size=8, num_steps=50000,
generate_every=5000, save_every=5000)
# ? send message to parent JS process
print(generate_message()) match sys.argv[1]:
sys.stdout.flush() case "generate":
# ? send message to parent JS process
print(generate_message())
sys.stdout.flush()
case "train":
train_ai()
case _:
raise NameError("Argument not provided.")