From 9f088473ac21970a5504d339c55a03db14d317bc Mon Sep 17 00:00:00 2001
From: BuildTools <unconfigured@null.spigotmc.org>
Date: Tue, 12 Jan 2021 20:36:52 +0100
Subject: [PATCH] Added training function + dir cleanup

---
 .gitignore |  8 ++++++--
 run.py     | 23 +++++++++++++++++------
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 48a69bb..e24d8a8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,9 @@
 # Exclude IDE Files
-.idea
+.idea/
 
 # Exclude VENV
-/venv/
\ No newline at end of file
+venv/
+
+# Exclude ChatLogs
+
+chatlogs/*
\ No newline at end of file
diff --git a/run.py b/run.py
index d023ac6..aa1464f 100644
--- a/run.py
+++ b/run.py
@@ -12,10 +12,6 @@ config = {"botname": "Sheepy",
           "corpus": ["chatterbot.corpus.english", "chatterbot.corpus.uswest"],
           "token": "NzU2NTgxNjQ0NTY3MTgzMzcw.X2T7kA.J7RthGGYMjgE_ZreLrGo2A3tCvg"}
 
-file = open("token", "r")
-token = file.read()
-file.close()
-
 bot = discord.Client()
 chatbot = ChatBot(config["botname"])
 
@@ -32,6 +28,7 @@ def printversion():
 
 async def download():
     await bot.change_presence(activity=discord.Game(name=" downloading Message History"))
+    os.system("rm *.tmp.yml")
     for trainchannel in config["trainchannels"]:
         channel = await bot.fetch_channel(trainchannel[0])
         counter = 0
@@ -40,10 +37,22 @@ async def download():
             messages[counter] = message.content
             counter += 1
         messages.reverse()
-        with open(str(trainchannel[0]) + ".yml", "w") as logfile:
+        with open("chatlogs/" + str(trainchannel[0]) + ".tmp.yml", "w") as logfile:
             yaml.dump(messages, logfile)
 
 
+async def train():
+    os.system("rm db.sqlite3")
+    corpustrainer = ChatterBotCorpusTrainer(chatbot)
+    listtrainer = ListTrainer(chatbot)
+    for corpus in config["corpus"]:
+        corpustrainer.train(corpus)
+    for log in os.listdir("chatlogs"):
+        with open(log, "r") as logfile:
+            log = yaml.load(logfile)
+        listtrainer.train(log)
+
+
 @bot.event
 async def on_ready():
     print("Logged in!")
@@ -65,6 +74,8 @@ async def on_ready():
             await bot.close()
     if tasks[0]:
         await download()
+    if tasks[1]:
+        await train()
 
 
-bot.run(token)
+bot.run(config["token"])