Added some more models

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)

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')
 
+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()