from chatterbot import ChatBot

if __name__ == '__main__':
    bot = ChatBot("bot", storage_adapter='chatterbot.storage.SQLStorageAdapter',
                  database_uri='sqlite:///database.sqlite3', read_only=True)

    flag = True
    while flag:
        msg = input("user:")
        if msg == '退出':
            flag = False
            break
        res = bot.get_response(msg)
        print(f'bot:{res}')

class CustomTrain(Trainer):
    def train(self, filePathList):
        print("开始自定义训练")
        for filePath in filePathList:
            with open(filePath, 'r') as f:
                previous_statement_text = None
                previous_statement_search_text = ''
                statements_to_create = []
                for line in f.readlines():
                    if line == '\n':
                        continue
                    if '%robotname%' in line:
                        line = line.replace('%robotname%', self.chatbot.name)
                    if 'XXX' in line:
                        line = line.replace('XXX', '你')

                    statement_search_text = self.chatbot.storage.tagger.get_bigram_pair_string(line)
                    statement = self.get_preprocessed_statement(
                        Statement(
                            text=line,
                            search_text=statement_search_text,
                            in_response_to=previous_statement_text,
                            search_in_response_to=previous_statement_search_text,
                            conversation='training'
                        )
                    )
                    previous_statement_text = statement.text
                    previous_statement_search_text = statement_search_text

                    statements_to_create.append(statement)
                self.chatbot.storage.create_many(statements_to_create)


if __name__ == '__main__':
    bot = ChatBot("001", storage_adapter='chatterbot.storage.SQLStorageAdapter',
                  database_uri='sqlite:///database.sqlite3')
    trainer = CustomTrain(bot)
    trainer.train([
        '/home/orange/ck.txt'
        '/home/orange/xhy.txt'
    ])