SheepBotFrame/run.py

87 lines
2.3 KiB
Python
Raw Normal View History

import os
import sys
import yaml
import discord
from chatterbot import ChatBot
from chatterbot.trainers import ChatterBotCorpusTrainer, ListTrainer
with open("config.yml", "r") as file:
config = yaml.load(file)
2021-01-12 12:49:03 -06:00
bot = discord.Client()
tasks = [False, False]
def printhelp():
print("Help lol")
def printversion():
print("Somer Version")
async def download():
await bot.change_presence(activity=discord.Game(name=" downloading Message History"))
2021-01-12 13:37:19 -06:00
os.system("rm chatlogs/*.tmp.yml")
for trainchannel in config["trainchannels"]:
channel = await bot.fetch_channel(trainchannel[0])
counter = 0
messages = [None] * trainchannel[1]
async for message in channel.history(limit=trainchannel[1]):
messages[counter] = message.content
counter += 1
messages.reverse()
2021-01-12 13:36:52 -06:00
with open("chatlogs/" + str(trainchannel[0]) + ".tmp.yml", "w") as logfile:
yaml.dump(messages, logfile)
2021-01-12 13:36:52 -06:00
async def train():
corpustrainer = ChatterBotCorpusTrainer(chatbot)
listtrainer = ListTrainer(chatbot)
for corpus in config["corpus"]:
corpustrainer.train(corpus)
for log in os.listdir("chatlogs"):
with open(f"chatlogs/{log}", "r") as logfile:
2021-01-12 13:36:52 -06:00
log = yaml.load(logfile)
listtrainer.train(log)
@bot.event
async def on_ready():
print("Logged in!")
args = sys.argv
args.pop(0)
for arg in args:
if arg == "--help" or arg == "-?":
printhelp()
await bot.close()
raise SystemExit(0)
elif arg == "--version" or arg == "-v":
printversion()
await bot.close()
raise SystemExit(0)
elif arg == "--download" or arg == "-d":
tasks[0] = True
elif arg == "--train" or arg == "-t":
tasks[1] = True
else:
print(f"Unexspected {arg}")
printhelp()
await bot.close()
raise SystemExit(1)
if tasks[0]:
await download()
print("Finished download!")
2021-01-12 13:36:52 -06:00
if tasks[1]:
2021-01-13 12:54:25 -06:00
os.system("rm db.sqlite3")
global chatbot
chatbot = ChatBot(config["botname"])
2021-01-12 13:36:52 -06:00
await train()
print("Finished train!")
2021-01-13 12:54:25 -06:00
else:
global chatbot
2021-01-13 12:54:25 -06:00
chatbot = ChatBot(config["botname"])
2021-01-12 12:49:03 -06:00
bot.run(config["token"])