Skip to content
Fetching contributors…
Cannot retrieve contributors at this time
executable file 251 lines (192 sloc) 8.57 KB
# Load the Ubuntu dialog corpus
# Available from here:
from multiprocessing import Pool
from pathlib import Path
from collections import OrderedDict
from urllib.request import urlretrieve
import os
import argparse
import tarfile
import pickle
from tqdm import tqdm
import pandas as pd
from model.utils import Tokenizer, Vocab, PAD_TOKEN, SOS_TOKEN, EOS_TOKEN
project_dir = Path(__file__).resolve().parent
datasets_dir = project_dir.joinpath('datasets/')
ubuntu_dir = datasets_dir.joinpath('ubuntu/')
ubuntu_meta_dir = ubuntu_dir.joinpath('meta/')
dialogs_dir = ubuntu_dir.joinpath('dialogs/')
# Tokenizer
tokenizer = Tokenizer('spacy')
def prepare_ubuntu_data():
"""Download and unpack dialogs"""
tar_filename = 'ubuntu_dialogs.tgz'
url = ''
tarfile_path = ubuntu_dir.joinpath(tar_filename)
metadata_url = ''
if not datasets_dir.exists():
if not ubuntu_dir.exists():
if not ubuntu_meta_dir.exists():
# Prepare Dialog data
if not dialogs_dir.joinpath("10/1.tst").exists():
# Download Dialog tarfile
if not tarfile_path.exists():
print(f"Downloading {url} to {tarfile_path}")
urlretrieve(url, tarfile_path)
print(f"Successfully downloaded {tarfile_path}")
# Unpack tarfile
if not dialogs_dir.exists():
print("Unpacking dialogs ... (This can take 5~10 mins.)")
with as tar:
print("Archive unpacked.")
# Download metadata
if not ubuntu_meta_dir.joinpath('trainfiles.csv').exists():
print('Downloading metadata ... (This can take 5~10 mins.)')
for filename in ['trainfiles.csv', 'valfiles.csv', 'testfiles.csv']:
csv_path = ubuntu_meta_dir.joinpath(filename)
print(f"Downloading {metadata_url+filename} to {csv_path}")
urlretrieve(metadata_url + filename, csv_path)
print(f"Successfully downloaded {csv_path}")
print('Ubuntu Data prepared!')
def get_dialog_path_list(dataset='train'):
if dataset == 'train':
filename = 'trainfiles.csv'
elif dataset == 'test':
filename = 'testfiles.csv'
elif dataset == 'valid':
filename = 'valfiles.csv'
with open(ubuntu_meta_dir.joinpath(filename)) as f:
dialog_path_list = []
for line in f:
file, dir = line.strip().split(",")
path = dialogs_dir.joinpath(dir, file)
return dialog_path_list
def read_and_tokenize(dialog_path, min_turn=3):
Read conversation
dialog_path (str): path of dialog (tsv format)
dialogs: (list of list of str) [dialog_length, sentence_length]
users: (list of str); [2]
with open(dialog_path, 'r', encoding='utf-8') as f:
# Go through the dialog
first_turn = True
dialog = []
users = []
same_user_utterances = [] # list of sentences of current user
for line in f:
_time, speaker, _listener, sentence = line.split('\t')
if first_turn:
last_speaker = speaker
first_turn = False
# Speaker has changed
if last_speaker != speaker:
same_user_utterances = []
last_speaker = speaker
# All users in conversation (len: 2)
users = list(OrderedDict.fromkeys(users))
# 1. Concatenate consecutive sentences of single user
# 2. Tokenize
dialog = [tokenizer(" ".join(sentence)) for sentence in dialog]
if len(dialog) < min_turn:
print(f"Dialog {dialog_path} length ({len(dialog)}) < minimum required length {min_turn}")
return []
return dialog #, users
def pad_sentences(conversations, max_sentence_length=30, max_conversation_length=10):
def pad_tokens(tokens, max_sentence_length=max_sentence_length):
n_valid_tokens = len(tokens)
if n_valid_tokens > max_sentence_length - 1:
tokens = tokens[:max_sentence_length - 1]
n_pad = max_sentence_length - n_valid_tokens - 1
tokens = tokens + [EOS_TOKEN] + [PAD_TOKEN] * n_pad
return tokens
def pad_conversation(conversation):
conversation = [pad_tokens(sentence) for sentence in conversation]
return conversation
all_padded_sentences = []
all_sentence_length = []
for conversation in conversations:
if len(conversation) > max_conversation_length:
conversation = conversation[:max_conversation_length]
sentence_length = [min(len(sentence) + 1, max_sentence_length) # +1 for EOS token
for sentence in conversation]
sentences = pad_conversation(conversation)
# [n_conversations, n_sentence (various), max_sentence_length]
sentences = all_padded_sentences
# [n_conversations, n_sentence (various)]
sentence_length = all_sentence_length
return sentences, sentence_length
if __name__ == '__main__':
parser = argparse.ArgumentParser()
# Maximum valid length of sentence
# => SOS/EOS will surround sentence (EOS for source / SOS for target)
# => maximum length of tensor = max_sentence_length + 1
parser.add_argument('-s', '--max_sentence_length', type=int, default=30)
parser.add_argument('-c', '--max_conversation_length', type=int, default=10)
# Vocabulary
parser.add_argument('--max_vocab_size', type=int, default=20000)
parser.add_argument('--min_vocab_frequency', type=int, default=5)
# Multiprocess
parser.add_argument('--n_workers', type=int, default=os.cpu_count())
args = parser.parse_args()
max_sent_len = args.max_sentence_length
max_conv_len = args.max_conversation_length
max_vocab_size = args.max_vocab_size
min_freq = args.min_vocab_frequency
n_workers = args.n_workers
min_turn = 3
# Download and unpack dialogs if necessary.
def to_pickle(obj, path):
with open(path, 'wb') as f:
pickle.dump(obj, f)
for split_type in ['train', 'test', 'valid']:
print(f'Processing {split_type} dataset...')
split_data_dir = ubuntu_dir.joinpath(split_type)
# List of dialogs (tsv)
dialog_path_list = get_dialog_path_list(split_type)
print(f'Tokenize.. (n_workers={n_workers})')
def _tokenize_conversation(dialog_path):
return read_and_tokenize(dialog_path)
with Pool(n_workers) as pool:
conversations = list(tqdm(pool.imap(_tokenize_conversation, dialog_path_list),
# Filter too short conversations
conversations = list(filter(lambda x: len(x) >= min_turn, conversations))
# conversations: padded_sentences
# [n_conversations, conversation_length (various), max_sentence_length]
# sentence_length: list of length of sentences
# [n_conversations, conversation_length (various)]
conversation_length = [min(len(conversation), max_conv_len)
for conversation in conversations]
sentences, sentence_length = pad_sentences(
print('Saving preprocessed data at', split_data_dir)
to_pickle(conversation_length, split_data_dir.joinpath('conversation_length.pkl'))
to_pickle(sentences, split_data_dir.joinpath('sentences.pkl'))
to_pickle(sentence_length, split_data_dir.joinpath('sentence_length.pkl'))
if split_type == 'train':
print('Save Vocabulary...')
vocab = Vocab(tokenizer)
vocab.update(max_size=max_vocab_size, min_freq=min_freq)
print('Vocabulary size: ', len(vocab))
vocab.pickle(ubuntu_dir.joinpath('word2id.pkl'), ubuntu_dir.joinpath('id2word.pkl'))
You can’t perform that action at this time.