这次的对联生成任务需要两个模型来完成，第一个模型根据首字来生成上联，例如：输入 “月” 预测出 “月似高人明大道”， 第二个模型则根据输入的上联预测出下联，例如：输入 “月似高人明大道” 预测出 “山如前辈傲雄关”。原始数据包括 train 和 test 目录，每个目录中分别包含 in.txt 和 out.txt 分别对应上联和下联，我们接下来就使用 train 目录中的对联数据进行训练。

train 数据共有 74 万数据，接下来，对数据进行部分处理。

1. 语料清洗

首先，对语料进行简单的处理，去除一些重复的数据，一些包含了 “妓女” 关键字的数据，以及一些包含逗号、问号的语料，例如：”历 史 名 城 ， 九 水 回 澜 ， 飞 扬 吴 楚 三 千 韵”、”万 险 何 辞 ？ 为 圆 大 漠 科 研 梦” 等等，此时数据量从 74 万下降到了 31.5 万左右。但是我们仍然不打算使用全部的数据，而是只使用其中的 2000 条数据进行训练。

完整的处理代码如下：

def clean_corpus():

    data_path1 = 'couplet/train/in.txt'
    data_path2 = 'couplet/train/out.txt'

    lines = []
    for line1, line2 in zip(open(data_path1), open(data_path2)):

        line1 = ''.join(line1.strip().split())
        line2 = ''.join(line2.strip().split())

        if line1.find('，') != -1 or line1.find('？') != -1:
            continue

        if line1.find('妓女') != -1 or line2.find('妓女') != -1:
            continue

        if len(line1) != len(line2):
            continue

        if len(line1) < 5:
            continue


        lines.append([line1, line2])

    # 根据上联去重
    lines = pd.DataFrame(data=lines, columns=['first', 'second'])
    lines = lines[~lines['first'].duplicated()]

    print('全部数据总量:', len(lines))

    # 存储训练数据
    lines = lines[:2000]
    lines.to_csv(train_path, index=False)

2. 构建词表

为了能够将文本数据进行数值化，我们需要构建词表。由于我们打算训练两个模型分对对应生成上联和下联，所以针对上联和下联分别构建词表，即：上联模型的词表、下联模型的词表，词表构建完毕之后，我们直接将其序列化磁盘中。其中，index_to_word 完成索引到字的映射，word_to_index 完成字到索引的映射。

构建上联词表代码：

vocab_single_path = 'data/vocab-single.pkl'


def build_single_vocab():

    # 读取数据
    train_data = pd.read_csv(train_path)
    train_data = train_data['first'].values.tolist()

    index_to_word = {}
    word_to_index = {}

    words = []
    for line in train_data:
        words.extend(line)

    words = list(set(words))
    words.insert(0, '[EOS]')

    for index, word in enumerate(words):
        index_to_word[index] = word
        word_to_index[word] = index

    # 存储词表
    save_data = {
        'index_to_word': index_to_word,
        'word_to_index': word_to_index,
        'vocab_size': len(index_to_word)
    }
    pickle.dump(save_data, open(vocab_single_path, 'wb'), protocol=3)

构建下联词表代码：

vocab_double_path = 'data/vocab-double.pkl'

def build_double_vocab():

    # 读取数据
    train_data = pd.read_csv(train_path)
    train_data = train_data[['first', 'second']].values.tolist()

    index_to_word = {}
    word_to_index = {}

    words = []
    for line1, line2 in train_data:
        words.extend(line1)
        words.extend(line2)

    words = list(set(words))
    words.insert(0, '[SOS]')
    words.insert(0, '[EOS]')

    for index, word in enumerate(words):
        index_to_word[index] = word
        word_to_index[word] = index

    save_data = {
        'index_to_word': index_to_word,
        'word_to_index': word_to_index,
        'vocab_size': len(index_to_word)
    }

    # 存储词表
    pickle.dump(save_data, open(vocab_double_path, 'wb'), protocol=3)

3. 训练语料

前面工作完成之后，我们将选中的 2000 条数据转换为在词表中的索引表示，便于后续送入到模型计算。处理完成之后，依然序列化到文件中。

上联语料代码：

vocab_single_path = 'data/vocab-single.pkl'


def build_single_train_corpus():

    # 训练数据
    train_data = pd.read_csv(train_path)
    train_data = train_data['first'].values.tolist()

    # 语料词表
    couple_single_vocab = pickle.load(open(vocab_single_path, 'rb'))
    word_to_index = couple_single_vocab['word_to_index']
    EOS = word_to_index['[EOS]']

    # 构建训练语料
    data = []
    for line in train_data:
        temp = []
        for word in line:
            temp.append(word_to_index[word])
        temp.append(EOS)
        data.append(temp)

    train_save_data = {
        'train_data': data,
        'train_size': len(train_data)
    }

    # 存储数据
    pickle.dump(train_save_data, open(train_single_path, 'wb'))

下联语料代码：

vocab_double_path = 'data/vocab-double.pkl'


def build_double_train_corpus():

    # 训练数据
    train_data = pd.read_csv(train_path)
    train_data = train_data[['first', 'second']].values.tolist()

    # 语料词表
    couple_single_vocab = pickle.load(open(vocab_double_path, 'rb'))
    word_to_index = couple_single_vocab['word_to_index']
    SOS = word_to_index['[SOS]']
    EOS = word_to_index['[EOS]']

    # 构建训练语料
    data = []
    for line1, line2 in train_data:
        temp1, temp2 = [], []
        for word1, word2 in zip(line1, line2):
            temp1.append(word_to_index[word1])
            temp2.append(word_to_index[word2])
        
        # temp1.append(SOS)
        temp2.insert(0, SOS)
        temp2.append(EOS)
        
        data.append([temp1, temp2])

    train_save_data = {
        'train_data': data,
        'train_size': len(train_data)
    }

    # 存储数据
    pickle.dump(train_save_data, open(train_double_path, 'wb'))

3. 执行入口

在执行完下面的代码之后，会在 data 目录下创建以下文件：

1. train.csv 从原始语料中选择出的 2000 条训练数据
2. train-double.pkl 用于对联下联训练的语料，该语料内容是基于 train.csv 内容
3. train-single.pkl 用于对联上联训练的语料，该语料内容是基于 train.csv 内容
4. vocab-double.pkl 用于对联下联的词表文件
5. vocab-single.pkl 用于对联上联的词表文件

def report():
    print('训练数据总量:', len(pd.read_csv(train_path)))
    print('上联词汇总量:', pickle.load(open(vocab_single_path, 'rb'))['vocab_size'])
    print('下联词汇总量:', pickle.load(open(vocab_double_path, 'rb'))['vocab_size'])


def start_preprocessing():

    # 1. 语料清洗
    clean_corpus()
    # 2. 构建词表
    build_single_vocab()
    build_double_vocab()
    # 3. 训练语料
    build_single_train_corpus()
    build_double_train_corpus()
    # 4. 基本信息
    report()


if __name__ == '__main__':
    start_preprocessing()