bobloy
4 years ago
committed by
GitHub
commit
ea126db0c5
@ -0,0 +1,73 @@
|
||||
from chatterbot.storage import StorageAdapter, SQLStorageAdapter
|
||||
|
||||
|
||||
class MyDumbSQLStorageAdapter(SQLStorageAdapter):
|
||||
def __init__(self, **kwargs):
|
||||
super(SQLStorageAdapter, self).__init__(**kwargs)
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
self.database_uri = kwargs.get("database_uri", False)
|
||||
|
||||
# None results in a sqlite in-memory database as the default
|
||||
if self.database_uri is None:
|
||||
self.database_uri = "sqlite://"
|
||||
|
||||
# Create a file database if the database is not a connection string
|
||||
if not self.database_uri:
|
||||
self.database_uri = "sqlite:///db.sqlite3"
|
||||
|
||||
self.engine = create_engine(
|
||||
self.database_uri, convert_unicode=True, connect_args={"check_same_thread": False}
|
||||
)
|
||||
|
||||
if self.database_uri.startswith("sqlite://"):
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy import event
|
||||
|
||||
@event.listens_for(Engine, "connect")
|
||||
def set_sqlite_pragma(dbapi_connection, connection_record):
|
||||
dbapi_connection.execute("PRAGMA journal_mode=WAL")
|
||||
dbapi_connection.execute("PRAGMA synchronous=NORMAL")
|
||||
|
||||
if not self.engine.dialect.has_table(self.engine, "Statement"):
|
||||
self.create_database()
|
||||
|
||||
self.Session = sessionmaker(bind=self.engine, expire_on_commit=True)
|
||||
|
||||
|
||||
class AsyncSQLStorageAdapter(SQLStorageAdapter):
|
||||
def __init__(self, **kwargs):
|
||||
super(SQLStorageAdapter, self).__init__(**kwargs)
|
||||
|
||||
self.database_uri = kwargs.get("database_uri", False)
|
||||
|
||||
# None results in a sqlite in-memory database as the default
|
||||
if self.database_uri is None:
|
||||
self.database_uri = "sqlite://"
|
||||
|
||||
# Create a file database if the database is not a connection string
|
||||
if not self.database_uri:
|
||||
self.database_uri = "sqlite:///db.sqlite3"
|
||||
|
||||
async def initialize(self):
|
||||
# from sqlalchemy import create_engine
|
||||
from aiomysql.sa import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
self.engine = await create_engine(self.database_uri, convert_unicode=True)
|
||||
|
||||
if self.database_uri.startswith("sqlite://"):
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy import event
|
||||
|
||||
@event.listens_for(Engine, "connect")
|
||||
def set_sqlite_pragma(dbapi_connection, connection_record):
|
||||
dbapi_connection.execute("PRAGMA journal_mode=WAL")
|
||||
dbapi_connection.execute("PRAGMA synchronous=NORMAL")
|
||||
|
||||
if not self.engine.dialect.has_table(self.engine, "Statement"):
|
||||
self.create_database()
|
||||
|
||||
self.Session = sessionmaker(bind=self.engine, expire_on_commit=True)
|
||||
@ -0,0 +1,351 @@
|
||||
import asyncio
|
||||
import csv
|
||||
import html
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
import time
|
||||
from functools import partial
|
||||
|
||||
from chatterbot import utils
|
||||
from chatterbot.conversation import Statement
|
||||
from chatterbot.tagging import PosLemmaTagger
|
||||
from chatterbot.trainers import Trainer
|
||||
from redbot.core.bot import Red
|
||||
from dateutil import parser as date_parser
|
||||
from redbot.core.utils import AsyncIter
|
||||
|
||||
log = logging.getLogger("red.fox_v3.chatter.trainers")
|
||||
|
||||
|
||||
class KaggleTrainer(Trainer):
|
||||
def __init__(self, chatbot, datapath: pathlib.Path, **kwargs):
|
||||
super().__init__(chatbot, **kwargs)
|
||||
|
||||
self.data_directory = datapath / kwargs.get("downloadpath", "kaggle_download")
|
||||
|
||||
self.kaggle_dataset = kwargs.get(
|
||||
"kaggle_dataset",
|
||||
"Cornell-University/movie-dialog-corpus",
|
||||
)
|
||||
|
||||
# Create the data directory if it does not already exist
|
||||
if not os.path.exists(self.data_directory):
|
||||
os.makedirs(self.data_directory)
|
||||
|
||||
def is_downloaded(self, file_path):
|
||||
"""
|
||||
Check if the data file is already downloaded.
|
||||
"""
|
||||
if os.path.exists(file_path):
|
||||
self.chatbot.logger.info("File is already downloaded")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def download(self, dataset):
|
||||
import kaggle # This triggers the API token check
|
||||
|
||||
future = await asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
partial(
|
||||
kaggle.api.dataset_download_files,
|
||||
dataset=dataset,
|
||||
path=self.data_directory,
|
||||
quiet=False,
|
||||
unzip=True,
|
||||
),
|
||||
)
|
||||
|
||||
def train(self, *args, **kwargs):
|
||||
log.error("See asynctrain instead")
|
||||
|
||||
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:
|
||||
# # log.info(f"Good line: {row}")
|
||||
# pass
|
||||
#
|
||||
# # lines_dict = {row[0].strip('"'): row[4] for row in reader_list}
|
||||
|
||||
statements_from_file = []
|
||||
save_every = 300
|
||||
count = 0
|
||||
|
||||
# [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)
|
||||
|
||||
count += 1
|
||||
if count >= save_every:
|
||||
if statements_from_file:
|
||||
self.chatbot.storage.create_many(statements_from_file)
|
||||
statements_from_file = []
|
||||
count = 0
|
||||
|
||||
if statements_from_file:
|
||||
self.chatbot.storage.create_many(statements_from_file)
|
||||
|
||||
log.info(f"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,
|
||||
)
|
||||
|
||||
async def asynctrain(self, *args, **kwargs):
|
||||
extracted_dir = self.data_directory / "Ubuntu-dialogue-corpus"
|
||||
|
||||
# Download and extract the Ubuntu dialog corpus if needed
|
||||
if not extracted_dir.exists():
|
||||
await self.download(self.kaggle_dataset)
|
||||
else:
|
||||
log.info("Ubuntu dialogue already downloaded")
|
||||
if not extracted_dir.exists():
|
||||
raise FileNotFoundError("Did not extract in the expected way")
|
||||
|
||||
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")
|
||||
|
||||
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()
|
||||
|
||||
tagger = PosLemmaTagger(language=self.chatbot.storage.tagger.language)
|
||||
|
||||
with open(extracted_dir / dialogue_file, "r", encoding="utf-8") as dg:
|
||||
reader = csv.DictReader(dg)
|
||||
|
||||
next(reader) # Skip the header
|
||||
|
||||
last_dialogue_id = None
|
||||
previous_statement_text = None
|
||||
previous_statement_search_text = ""
|
||||
statements_from_file = []
|
||||
|
||||
save_every = 50
|
||||
count = 0
|
||||
|
||||
async for row in AsyncIter(reader):
|
||||
dialogue_id = row["dialogueID"]
|
||||
if dialogue_id != last_dialogue_id:
|
||||
previous_statement_text = None
|
||||
previous_statement_search_text = ""
|
||||
last_dialogue_id = dialogue_id
|
||||
count += 1
|
||||
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:
|
||||
statement = Statement(
|
||||
text=row["text"],
|
||||
in_response_to=previous_statement_text,
|
||||
conversation="training",
|
||||
# created_at=date_parser.parse(row["date"]),
|
||||
persona=row["from"],
|
||||
)
|
||||
|
||||
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:
|
||||
self.chatbot.storage.create_many(statements_from_file)
|
||||
|
||||
log.info(f"Training took {time.time() - start_time} seconds.")
|
||||
|
||||
|
||||
class TwitterCorpusTrainer(Trainer):
|
||||
pass
|
||||
# def train(self, *args, **kwargs):
|
||||
# """
|
||||
# Train the chat bot based on the provided list of
|
||||
# statements that represents a single conversation.
|
||||
# """
|
||||
# import twint
|
||||
#
|
||||
# c = twint.Config()
|
||||
# c.__dict__.update(kwargs)
|
||||
# twint.run.Search(c)
|
||||
#
|
||||
#
|
||||
# previous_statement_text = None
|
||||
# previous_statement_search_text = ''
|
||||
#
|
||||
# statements_to_create = []
|
||||
#
|
||||
# for conversation_count, text in enumerate(conversation):
|
||||
# if self.show_training_progress:
|
||||
# utils.print_progress_bar(
|
||||
# 'List Trainer',
|
||||
# conversation_count + 1, len(conversation)
|
||||
# )
|
||||
#
|
||||
# statement_search_text = self.chatbot.storage.tagger.get_text_index_string(text)
|
||||
#
|
||||
# statement = self.get_preprocessed_statement(
|
||||
# Statement(
|
||||
# text=text,
|
||||
# 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)
|
||||
Loading…
Reference in new issue