tensor 增加一维度_PyTorch搭建聊天机器人（一）词表与数据加载器

国庆无聊逛了逛PyTorch的tutorial，其中有一篇chatbot的搭建蛮有意思的。

Chatbot Tutorial​pytorch.org

作为一个蒟蒻大二，我看了看tutorial涉及到的论文，并且自己按照batch_first=True动手写了写，算是有点收获吧。打算写三四篇文章总结一下技术细节。顺序大概是：数据加载器、网络前向逻辑、训练逻辑、评估逻辑。

使用pycharm编写项目，代码分为四个文件：process.py、neural_network.py、train.py、evaluate.py。

先大致说一下搭建chatbot的思路吧，其实很简单：这里的chatbot是基于带Luong attention机制的seq2seq。研究过NLP的同学应该对seq2seq很熟悉，它可以将任意长度的时序信息映射到任意长度，在基于深度神经网络的机器翻译中使用广泛。

实际上，中文翻译成英文就是训练出一个中文序列到英文序列的映射，而我们的chatbot不就是一个句子到句子的映射吗？在不考虑上下语境的情况下，聊天机器人可以使用seq2seq搭建。如此搭建的聊天机器人对用户输入的语句给出的回复更像是将用户说的话翻译成了用户希望得到的回复。那么假设我们已经对seq2seq很熟悉了，那么只需要使用一条条对话（下面叫dialog或者pair）作为数据，训练这个seq2seq模型就可以得到这个训练集风格的chatbot了。

tutorial使用的数据是Cornell Movie-Dialogs，下载地址。这部分数据的编码格式不是utf-8，如果你对编码转换这部分不感兴趣，可以直接使用笔者仓库中./data中的tsv数据。后面的程序中将会直接使用tsv数据。

笔者仓库链接如下：

LSTM-Kirigaya/chatbot-based-on-seq2seq2​github.com

提前说明一下，对话数据集中的每个pair中，我们把第一句话成为input_dialog，后面一句回复的话称为ouput_dialog。

---

下面完成process.py，这个文件完成词表建立和数据加载器的建立。

说明：下面所有数据组织都是按照batch_first来的，也就是所有torch张量的第一个维度是batch_size

先引入需要的库

from itertools import zip_longest
import random
import torch

构建词表

第一步我们需要构建词表，因为网络中只会传递张量，我们需要通过构建词表将每个单词映射成一个个单词索引（后面成为index），也就是将一句话转化为index序列。

词表中最核心的数据是三个python类型的词典：

• word2index：单词到其对应的index的映射。
• index2word：index到其对应的单词的映射。
• word2count：单词到其在数据集中的总数的映射。

构建词表的逻辑也很简单，只需要遍历数据集，每遇到一个词表中没有的单词，就根据已经添加单词的总数给与这个新的单词一个index，并由此给word2index和index2word两个字典增加新的元素。

程序如下：

# 用来构造字典的类
class vocab(object):
    def __init__(self, name, pad_token, sos_token, eos_token, unk_token):
        self.name = name
        self.pad_token = pad_token
        self.sos_token = sos_token
        self.eos_token = eos_token
        self.unk_token = unk_token
        self.trimmed = False                         # 代表这个词表对象是否经过了剪枝操作
        self.word2index = {"PAD" : pad_token, "SOS" : sos_token, "EOS" : eos_token, "UNK" : unk_token}
        self.word2count = {"UNK" : 0}
        self.index2word = {pad_token : "PAD", sos_token : "SOS", eos_token : "EOS", unk_token : "UNK"}
        self.num_words = 4  # 刚开始的四个占位符 pad(0), sos(1), eos(2)，unk(3) 代表目前遇到的不同的单词数量

    # 向voc中添加一个单词的逻辑
    def addWord(self, word):
        if word not in self.word2index:
            self.word2index[word] = self.num_words
            self.word2count[word] = 1
            self.index2word[self.num_words] = word
            self.num_words += 1
        else:
            self.word2count[word] += 1

    # 向voc中添加一个句子的逻辑
    def addSentence(self, sentence):
        for word in sentence.split():
            self.addWord(word)

    # 将词典中词频过低的单词替换为unk_token
    # 需要一个代表修剪阈值的参数min_count，词频低于这个参数的单词会被替换为unk_token，相应的词典变量也会做出相应的改变
    def trim(self, min_count):
        if self.trimmed:   # 如果已经裁剪过了，那就直接返回
            return
        self.trimmed = True

        keep_words = []
        keep_num = 0
        for word, count in self.word2count.items():
            if count >= min_count:
                keep_num += 1
                # 由于后面是通过对keep_word列表中的数据逐一统计，所以需要对count>1的单词重复填入
                for _ in range(count):
                    keep_words.append(word)

        print("keep words: {} / {} = {:.4f}".format(
           keep_num, self.num_words - 4, keep_num / (self.num_words - 4)
        ))

        # 重构词表
        self.word2index  = {"PAD" : self.pad_token, "SOS" : self.sos_token, "EOS" : self.eos_token, "UNK" : self.unk_token}
        self.word2count = {}
        self.index2word = {self.pad_token : "PAD", self.sos_token : "SOS", self.eos_token : "EOS", self.unk_token : "UNK"}
        self.num_words = 4

        for word in keep_words:
            self.addWord(word)

    # 读入数据，统计词频，并返回数据
    def load_data(self, path):
        pairs = []
        for line in open(path, "r", encoding="utf-8"):
            try:
                input_dialog, output_dialog = line.strip().split("t")
                self.addSentence(input_dialog.strip())
                self.addSentence(output_dialog.strip())
                pairs.append([input_dialog, output_dialog])
            except:
                pass
        return pairs

这个词表类，需要五个参数初始化：name、pad_token、sos_token、eos_token、unk_token。分别为词表的名称、填充词的index、句子开头标识符的index、句子结束标识符的index和未识别单词的index。

主要方法说明如下：

• __init__：完成词表的初始化。
• trim：根据min_count对词表进行剪枝。
• load_data：载入外部tsv数据，完成三个字典的搭建，并返回处理好的pairs。

处理input_dialog和output_dialog

有了词表，我们就可以根据词表把一句话转换成index序列，为此我们通过sentenceToIndex函数完成sentence到index sequence的转换，需要说明的是，为了让后续搭建的网络知道一句话已经结束了，我们需要给每个转换成的index序列的句子添加一个eos_token作为后缀：

# 将一句话转换成id序列(str->list)，结尾加上EOS
def sentenceToIndex(sentence, voc):
    return [voc.word2index[word] for word in sentence.split()] + [voc.eos_token]

接下来我们需要分别处理input_dialog和output_dialog。

处理input_dialog需要一个batchInput2paddedTensor函数，这个函数接受batch_size句没有处理过的input_dialog文字、将它们转换为index序列、填充pad_token、转换成batch_first=True的torch张量，返回处理好的torch张量和每句话的长度信息。

大致过程还是画张图吧。。。

代码如下：

# 将一个batch中的input_dialog转化为有pad填充的tensor，并返回tensor和记录长度的变量
# 返回的tensor是batch_first的
def batchInput2paddedTensor(batch, voc):
    # 先转换为id序列，但是这个id序列不对齐
    batch_index_seqs = [sentenceToIndex(sentence, voc) for sentence in batch]
    length_tensor = torch.tensor([len(index_seq) for index_seq in batch_index_seqs])
    # 下面填充0(PAD)，使得这个batch中的序列对齐
    zipped_list = list(zip_longest(*batch_index_seqs, fillvalue=voc.pad_token))
    padded_tensor = torch.tensor(zipped_list).t()
    return padded_tensor, length_tensor

处理output_dialog与input_dialog差不多，只不过需要多返回一个mask矩阵，所谓mask矩阵，就是将padded_tensor转换成bool类型。这些返回的量在后续的训练中都会使用到。output_dialog的处理如下：

# 将一个batch中的output_dialog转化为有pad填充的tensor，并返回tensor、mask和最大句长
# 返回的tensor是batch_first的
def batchOutput2paddedTensor(batch, voc):
    # 先转换为id序列，但是这个id序列不对齐
    batch_index_seqs = [sentenceToIndex(sentence, voc) for sentence in batch]
    max_length = max([len(index_seq) for index_seq in batch_index_seqs])
    # 下面填充0(PAD)，使得这个batch中的序列对齐
    zipped_list = list(zip_longest(*batch_index_seqs, fillvalue=voc.pad_token))
    padded_tensor = torch.tensor(zipped_list).t()
    # 得到padded_tensor对应的mask
    mask = torch.BoolTensor(zipped_list).t()
    return padded_tensor, mask, max_length

有了处理pair的函数，我们可以把上面的函数整合成一个数据加载器loader。数据加载器在深度学习中很重要，我们在训练中需要能够不重复的、快速地获取一个batch的格式化数据，这就是loader的功能，惬意舒适（in my dream...）的训练中，一个设计合理而高效的loader是必不可少的。

此处不再解释Python中生成器的概念，为了更加节省内存空间，我们将loader做成一个生成器：

# 获取数据加载器的函数
# 将输入的一个batch的dialog转换成id序列，填充pad，并返回训练可用的id张量和mask
def DataLoader(pairs, voc, batch_size, shuffle=True):
    if shuffle:
        random.shuffle(pairs)

    batch = []
    for idx, pair in enumerate(pairs):
        batch.append([pair[0], pair[1]])

        # 数据数量到达batch_size就yield出去并清空
        if len(batch) == batch_size:
            # 为了后续的pack_padded_sequence操作，我们需要给这个batch中的数据按照input_dialog的长度排序(降序)
            batch.sort(key=lambda x : len(x[0].split()), reverse=True)
            input_dialog_batch = []
            output_dialog_batch = []
            for pair in batch:
                input_dialog_batch.append(pair[0])
                output_dialog_batch.append(pair[1])

            input_tensor, input_length_tensor = batchInput2paddedTensor(input_dialog_batch, voc)
            output_tensor, mask, max_length = batchOutput2paddedTensor(output_dialog_batch, voc)

            # 清空临时缓冲区
            batch = []

            yield [
                input_tensor, input_length_tensor, output_tensor, mask, max_length
            ]

要写的函数差不多写好了，我们可以测试一下：

if __name__ == "__main__":
    PAD_token = 0  # 补足句长的pad占位符的index
    SOS_token = 1  # 代表一句话开头的占位符的index
    EOS_token = 2  # 代表一句话结尾的占位符的index
    UNK_token = 3  # 代表不在词典中的字符
    BATCH_SIZE = 64  # 一个batch中的对话数量（样本数量）
    MAX_LENGTH = 20  # 一个对话中每句话的最大句长
    MIN_COUNT = 3  # trim方法的修剪阈值

    # 实例化词表
    voc = vocab(name="corpus", pad_token=PAD_token, sos_token=SOS_token, eos_token=EOS_token, unk_token=UNK_token)
    # 为词表载入数据，统计词频，并得到对话数据
    pairs = voc.load_data(path="./data/dialog.tsv")
    print("total number of dialogs:", len(pairs))

    # 修剪与替换
    pairs = trimAndReplace(voc, pairs, MIN_COUNT)

    # 获取loader
    loader = DataLoader(pairs, voc, batch_size=5)

    batch_item_names = ["input_tensor", "input_length_tensor", "output_tensor", "mask", "max_length"]
    for batch_index, batch in enumerate(loader):
        for name, item in zip(batch_item_names, batch):
            print(f"n{name} : {item}")
        break

out:

total number of dialogs: 64223
keep words: 7821 / 17999 = 0.4345
Trimmed from 64223 pairs to 58362, 0.9087 of total

input_tensor : tensor([[ 123,   51,   48,    8,  918, 2227,  330, 3068,    7,    2],
        [ 302,  303,  102,   38,    3,   71,    3,    7,    2,    0],
        [ 158,    3,    7,    2,    0,    0,    0,    0,    0,    0],
        [ 188,    7,    2,    0,    0,    0,    0,    0,    0,    0],
        [ 563,    5,    2,    0,    0,    0,    0,    0,    0,    0]])

input_length_tensor : tensor([10,  9,  4,  3,  3])

output_tensor : tensor([[3244,    5,    2,    0,    0,    0,    0,    0,    0,    0],
        [  35,   37,   38,   68,   77,    5,    2,    0,    0,    0],
        [ 181,    5, 1233,   13, 1233,   13, 1222,    5,    2,    0],
        [ 102,   38,   45,  188,   99,  680, 1375,    5,    2,    0],
        [  26,  198,  118,   25,   51,   41,   48, 1597,    5,    2]])

mask : tensor([[ True,  True,  True, False, False, False, False, False, False, False],
        [ True,  True,  True,  True,  True,  True,  True, False, False, False],
        [ True,  True,  True,  True,  True,  True,  True,  True,  True, False],
        [ True,  True,  True,  True,  True,  True,  True,  True,  True, False],
        [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]])

max_length : 10

做好了数据加载器，后面就可以开始构建网络结构了。