Split models

create_models.py

@@ -1,6 +1,7 @@
 import os
 from dotenv import load_dotenv
 from telethon.sync import TelegramClient
+import markovify
 
 load_dotenv()
 
@@ -19,24 +20,23 @@ def cleanup(msg):
     '__'
   ]
 
-  try:
   for unwanted in unwanted_strings:
     msg = msg.replace(unwanted, '')
 
   return msg
-  except TypeError:
-    return msg
 
 
 with TelegramClient('bot', api_id, api_hash) as client:
-  f = open("data/data.txt", "w")
-
   for chat in chats:
     print(f'Collecting {chat}')
+    data = ''
     for message in client.iter_messages(chat):
-      try:
-        f.write(cleanup(message.text))
-      except:
-        pass
+      if (message.text):
+        data = data + cleanup(message.text)
 
+    text_model = markovify.Text(data)
+    # text_model.compile(inplace = True)
+    model_json = text_model.to_json()
+    f = open(f"data/{chat}.json", "w")
+    f.write(model_json)
     f.close()

main.py

@@ -1,4 +1,5 @@
 import os
+import random
 from dotenv import load_dotenv
 from telethon.sync import TelegramClient
 import markovify
@@ -15,20 +16,29 @@ api_hash = os.getenv('API_HASH')
 entity = os.getenv('ENTITY')
 dry_run = args.dry
 
-with open("data/data.txt", "r") as f:
-  text = f.read()
+def create_message(model):
+  message = ''
 
-text_model = markovify.Text(text)
-message = ''
-
-for i in range(5):
+  for i in range(5):
     try:
-    message = message + text_model.make_sentence()
+      message = message + model.make_sentence()
     except TypeError:
       pass
 
+  return message
+
+avalible_files = os.listdir('data')
+avalible_files.remove('.gitkeep')
+
+filename = random.choice(avalible_files)
+print(f'chosen {filename}')
+with open(os.path.join('data', filename), 'r') as f:
+  model_json = f.read()
+
+model = markovify.Text.from_json(model_json)
+
 if (dry_run):
-  print(message)
+  print(create_message(model))
 else:
   with TelegramClient('bot', api_id, api_hash) as client:
-    client.send_message(entity=entity, message=message)
+    client.send_message(entity=entity, message=create_message(model))