Add new kaggle trainers

chatter/trainers.py

@@ -1,5 +1,6 @@
 import asyncio
 import csv
+import html
 import logging
 import os
 import pathlib
@@ -56,13 +57,159 @@ class KaggleTrainer(Trainer):
             ),
         )
 
+    def train(self, *args, **kwargs):
+        log.error("See asynctrain instead")
 
-class UbuntuCorpusTrainer2(KaggleTrainer):
+    def asynctrain(self, *args, **kwargs):
+        raise self.TrainerInitializationException()
+
+
+class SouthParkTrainer(KaggleTrainer):
     def __init__(self, chatbot, datapath: pathlib.Path, **kwargs):
         super().__init__(
             chatbot,
             datapath,
             downloadpath="ubuntu_data_v2",
+            kaggle_dataset="tovarischsukhov/southparklines",
+            **kwargs,
+        )
+
+
+class MovieTrainer(KaggleTrainer):
+    def __init__(self, chatbot, datapath: pathlib.Path, **kwargs):
+        super().__init__(
+            chatbot,
+            datapath,
+            downloadpath="kaggle_movies",
+            kaggle_dataset="Cornell-University/movie-dialog-corpus",
+            **kwargs,
+        )
+
+    async def run_movie_training(self):
+        dialogue_file = "movie_lines.tsv"
+        conversation_file = "movie_conversations.tsv"
+        log.info(f"Beginning dialogue training on {dialogue_file}")
+        start_time = time.time()
+
+        tagger = PosLemmaTagger(language=self.chatbot.storage.tagger.language)
+
+        # [lineID, characterID, movieID, character name, text of utterance]
+        # File parsing from https://www.kaggle.com/mushaya/conversation-chatbot
+
+        with open(self.data_directory / conversation_file, "r", encoding="utf-8-sig") as conv_tsv:
+            conv_lines = conv_tsv.readlines()
+        with open(self.data_directory / dialogue_file, "r", encoding="utf-8-sig") as lines_tsv:
+            dialog_lines = lines_tsv.readlines()
+
+        # trans_dict = str.maketrans({"<u>": "__", "</u>": "__", '""': '"'})
+
+        lines_dict = {}
+        for line in dialog_lines:
+            _line = line[:-1].strip('"').split("\t")
+            if len(_line) >= 5:  # Only good lines
+                lines_dict[_line[0]] = (
+                    html.unescape(("".join(_line[4:])).strip())
+                    .replace("<u>", "__")
+                    .replace("</u>", "__")
+                    .replace('""', '"')
+                )
+            else:
+                log.debug(f"Bad line {_line}")
+
+        # collecting line ids for each conversation
+        conv = []
+        for line in conv_lines[:-1]:
+            _line = line[:-1].split("\t")[-1][1:-1].replace("'", "").replace(" ", ",")
+            conv.append(_line.split(","))
+
+        # conversations = csv.reader(conv_tsv, delimiter="\t")
+        #
+        # reader = csv.reader(lines_tsv, delimiter="\t")
+        #
+        #
+        #
+        # lines_dict = {}
+        # for row in reader:
+        #     try:
+        #         lines_dict[row[0].strip('"')] = row[4]
+        #     except:
+        #         log.exception(f"Bad line: {row}")
+        #         pass
+        #     else:
+        #         # print(f"Good line: {row}")
+        #         pass
+        #
+        # # lines_dict = {row[0].strip('"'): row[4] for row in reader_list}
+
+        statements_from_file = []
+
+        # [characterID of first, characterID of second, movieID, list of utterances]
+        async for lines in AsyncIter(conv):
+            previous_statement_text = None
+            previous_statement_search_text = ""
+
+            for line in lines:
+                text = lines_dict[line]
+                statement = Statement(
+                    text=text,
+                    in_response_to=previous_statement_text,
+                    conversation="training",
+                )
+
+                for preprocessor in self.chatbot.preprocessors:
+                    statement = preprocessor(statement)
+
+                statement.search_text = tagger.get_text_index_string(statement.text)
+                statement.search_in_response_to = previous_statement_search_text
+
+                previous_statement_text = statement.text
+                previous_statement_search_text = statement.search_text
+
+                statements_from_file.append(statement)
+
+            if statements_from_file:
+                print(statements_from_file)
+                self.chatbot.storage.create_many(statements_from_file)
+                statements_from_file = []
+
+        print("Training took", time.time() - start_time, "seconds.")
+
+    async def asynctrain(self, *args, **kwargs):
+        extracted_lines = self.data_directory / "movie_lines.tsv"
+        extracted_lines: pathlib.Path
+
+        # Download and extract the Ubuntu dialog corpus if needed
+        if not extracted_lines.exists():
+            await self.download(self.kaggle_dataset)
+        else:
+            log.info("Movie dialog already downloaded")
+        if not extracted_lines.exists():
+            raise FileNotFoundError(f"{extracted_lines}")
+
+        await self.run_movie_training()
+
+        return True
+
+        # train_dialogue = kwargs.get("train_dialogue", True)
+        # train_196_dialogue = kwargs.get("train_196", False)
+        # train_301_dialogue = kwargs.get("train_301", False)
+        #
+        # if train_dialogue:
+        #     await self.run_dialogue_training(extracted_dir, "dialogueText.csv")
+        #
+        # if train_196_dialogue:
+        #     await self.run_dialogue_training(extracted_dir, "dialogueText_196.csv")
+        #
+        # if train_301_dialogue:
+        #     await self.run_dialogue_training(extracted_dir, "dialogueText_301.csv")
+
+
+class UbuntuCorpusTrainer2(KaggleTrainer):
+    def __init__(self, chatbot, datapath: pathlib.Path, **kwargs):
+        super().__init__(
+            chatbot,
+            datapath,
+            downloadpath="kaggle_ubuntu",
             kaggle_dataset="rtatman/ubuntu-dialogue-corpus",
             **kwargs,
         )
@@ -91,6 +238,8 @@ class UbuntuCorpusTrainer2(KaggleTrainer):
         if train_301_dialogue:
             await self.run_dialogue_training(extracted_dir, "dialogueText_301.csv")
 
+        return True
+
     async def run_dialogue_training(self, extracted_dir, dialogue_file):
         log.info(f"Beginning dialogue training on {dialogue_file}")
         start_time = time.time()
@@ -120,6 +269,7 @@ class UbuntuCorpusTrainer2(KaggleTrainer):
                     if count >= save_every:
                         if statements_from_file:
                             self.chatbot.storage.create_many(statements_from_file)
+                            statements_from_file = []
                         count = 0
 
                 if len(row) > 0:
@@ -147,9 +297,6 @@ class UbuntuCorpusTrainer2(KaggleTrainer):
 
         print("Training took", time.time() - start_time, "seconds.")
 
-    def train(self, *args, **kwargs):
-        log.error("See asynctrain instead")
-
 
 class TwitterCorpusTrainer(Trainer):
     pass