diff --git a/create_models.py b/create_models.py index c606325..4277224 100644 --- a/create_models.py +++ b/create_models.py @@ -9,7 +9,10 @@ api_id = os.getenv('API_ID') api_hash = os.getenv('API_HASH') chats = [ 'tg1337const', - 'cyberboec' + 'cyberboec', + 'topor', + 'workBrothers', + 'anekdot_bb' ] def cleanup(msg): @@ -35,7 +38,7 @@ with TelegramClient('bot', api_id, api_hash) as client: data = data + cleanup(message.text) text_model = markovify.Text(data) - # text_model.compile(inplace = True) + text_model.compile(inplace = True) model_json = text_model.to_json() f = open(f"data/{chat}.json", "w") f.write(model_json) diff --git a/main.py b/main.py index fb6d4b1..020e3ee 100644 --- a/main.py +++ b/main.py @@ -7,6 +7,7 @@ from argparse import ArgumentParser parser = ArgumentParser() parser.add_argument("-d", "--dry", action='store_true', help="Run without sending a message") +parser.add_argument("-m", "--model", help="Specify model file") args = parser.parse_args() load_dotenv() @@ -15,6 +16,7 @@ api_id = os.getenv('API_ID') api_hash = os.getenv('API_HASH') entity = os.getenv('ENTITY') dry_run = args.dry +filename = args.model def create_message(model): message = '' @@ -30,7 +32,9 @@ def create_message(model): avalible_files = os.listdir('data') avalible_files.remove('.gitkeep') -filename = random.choice(avalible_files) +if (not filename): + filename = random.choice(avalible_files) + print(f'chosen {filename}') with open(os.path.join('data', filename), 'r') as f: model_json = f.read()