Google BERT 中文应用之《红楼梦》中对话人物提取

点击上方“迈微AI研习社”，选择“星标★”公众号

转载自：PaperWeekly 庞龙刚

之前看到过一篇文章，通过提取文章中对话的人物，分析人物之间的关系，很好奇如何通过编程的方式知道一句话是谁说的。但是遍搜网络没有发现类似的研究。

前段时间看到一个微信里的读书小程序，将人物对话都提取出来，将一本书的内容通过微信对话的方式表达出来，通过将对话的主角替换成读者的微信号以及用户头像，从而增加读者的代入感。试了之后非常佩服程序作者的巧思。这使得我写一个自然语言处理程序，提取书中对话，以及对话人物的念头更加强烈。

之前并没有多少 nlp 的经验，只零碎试过用 LSTM 训练写唐诗，用jieba做分词，用Google的 gensim 在WikiPedia中文语料上训练词向量。最近Google的BERT模型很火，运行了BERT的SQUAD阅读理解与问答系统，分类器以及特征提取例子之后，觉得这个任务可以用BERT微调来完成，在这里记录实验的粗略步骤，与君共勉。

我把训练数据和准备数据的脚本开源，放在gitlab上，开放下载。

LongGang Pang / SpeakerExtractiongitlab.com

链接：https://gitlab.com/snowhitiger/speakerextraction

BERT/SQUAD 预言的结果可以从 res.txt 里面找到。

（一）准备训练数据

《红楼梦》中的对话很好提取，大部分对话都有特定的格式，即一段话从：“开始，从”结束。使用 python 的正则表达式，可以很容易提取所有满足这样条件的对话。如果假设说出这段话的人的名字出现在这段话的前面，那么可以用这段话前面的一段话作为包含说话人(speaker ）的上下文（context）。如果说话人不存在这段上下文中，标签为空字符串。下面是第一步提取出的数据示例，

{'istart': 414, 'iend': 457, 'talk': '原来如此，下愚不知．但那宝玉既有如此的来历，又何以情迷至此，复又豁悟如此？还要请教。', 'context': '雨村听了，虽不能全然明白，却也十知四五，便点头叹道：'},
{'istart': 463, 'iend': 526, 'talk': '此事说来，老先生未必尽解．太虚幻境即是真如福地．一番阅册，原始要终之道，历历生平，如何不悟？仙草归真，焉有通灵不复原之理呢！', 'context': '士隐笑道：'},
{'istart': 552, 'iend': 588, 'talk': '宝玉之事既得闻命，但是敝族闺秀如此之多，何元妃以下算来结局俱属平常呢？', 'context': '雨村听着，却不明白了．知仙机也不便更问，因又说道：'},
{'istart': 880, 'iend': 891, 'talk': '此系后事，未便预说。', 'context': '士隐微微笑道：'},
{'istart': 19, 'iend': 45, 'talk': '老先生草庵暂歇，我还有一段俗缘未了，正当今日完结。', 'context': '食毕，雨村还要问自己的终身，士隐便道：'},
{'istart': 52, 'iend': 68, 'talk': '仙长纯修若此，不知尚有何俗缘？', 'context': '雨村惊讶道：'},
{'istart': 51, 'iend': 77, 'talk': '大士，真人，恭喜，贺喜！情缘完结，都交割清楚了么？', 'context': '这士隐自去度脱了香菱，送到太虚幻境，交那警幻仙子对册，刚过牌坊，见那一僧一道，缥渺而来．士隐接着说道：'},
{'istart': 75, 'iend': 243, 'talk': '我从前见石兄这段奇文，原说可以闻世传奇，所以曾经抄录，但未见返本还原．不知何时复有此一佳话，方知石兄下凡一次，磨出光明，修成圆觉，也可谓无复遗憾了．只怕年深日久，字迹模糊，反有舛错，不如我再抄录一番，寻个世上无事的人，托他传遍，知道奇而不奇，俗而不俗，真而不真，假而不假．或者尘梦劳人，聊倩鸟呼归去，山灵好客，更从石化飞来，亦未可知。', 'context': '这一日空空道人又从青埂峰前经过，见那补天未用之石仍在那里，上面字迹依然如旧，又从头的细细看了一遍，见后面偈文后又历叙了多少收缘结果的话头，便点头叹道：'},

大部分数据的上下文都很简单，比如 '士隐笑道：'等，但也有比较复杂的语境，比如 '这一日空空道人又从青埂峰前经过，见那补天未用之石仍在那里，上面字迹依然如旧，又从头的细细看了一遍，见后面偈文后又历叙了多少收缘结果的话头，便点头叹道：'。

（二）手动标记数据

为了训练机器，让它知道我想让它干什么，必须手动标记一些数据。我在Jupyter notebook 下写了一个简单的GUI程序，将每段话变成按钮，只需要点击需要标记数据的句首和句尾，程序会自动计算标记数据在上下文中的位置，并将记录保存到文本中。花了两个多小时，标记了大约1500多个数据，这些数据的最后几个例子如下，

{'uid': 1552, 'context': '黛玉又道：', 'speaker': '黛玉', 'istart': 0, 'iend': 2}
{'uid': 1553, 'context': '因念云：', 'speaker': None, 'istart': -1, 'iend': 0}
{'uid': 1554, 'context': '宝钗道：', 'speaker': '宝钗', 'istart': 0, 'iend': 2}
{'uid': 1555, 'context': '五祖便将衣钵传他．今儿这偈语，亦同此意了．只是方才这句机锋，尚未完全了结，这便丢开手不成？"黛玉笑道：', 'speaker': '黛玉', 'istart': 46, 'iend': 48}
{'uid': 1556, 'context': '宝玉自己以为觉悟，不想忽被黛玉一问，便不能答，宝钗又比出"语录"来，此皆素不见他们能者．自己想了一想：', 'speaker': '宝玉', 'istart': 0, 'iend': 2}
{'uid': 1557, 'context': '想毕，便笑道：', 'speaker': None, 'istart': -1, 'iend': 0}
{'uid': 1558, 'context': '说着，四人仍复如旧．忽然人报，娘娘差人送出一个灯谜儿，命你们大家去猜，猜着了每人也作一个进去．四人听说忙出去，至贾母上房．只见一个小太监，拿了一盏四角平头白纱灯，专为灯谜而制，上面已有一个，众人都争看乱猜．小太监又下谕道：', 'speaker': '小太监', 'istart': 103, 'iend': 106}
{'uid': 1559, 'context': '太监去了，至晚出来传谕：', 'speaker': '太监', 'istart': 0, 'iend': 2}

1500 个数据太少了，为了增加数据量，我又做了 data augmentation，将1500多个 speaker 插入到 1500 多个语境中，凭空生成了 200多万对训练数据。所以在训练数据中，有一些非常搞笑的内容，比如：

说毕走来，只见宝玉拄着拐棍，在当地骂袭人：

这个训练例子中的宝玉，原文应该是李嬷嬷。

（三）训练过程

简单构造 SQUAD 的中文训练和测试数据，训练并预测，结果输出在predictions.json中。

训练数据的 json 格式如下：

{"data" : [{"title": "红楼梦", "paragraphs":[{context and qas item 1}, {context and qas item 2}, ... {context and qas item i}, ..., {context and qas item n}]},
{"title": "寻秦记", "paragraphs":[{}, {}, {}]},
{"title": "xxxxxx", "paragraphs":[{}, {}, {}]}],
"version" : "speaker1.0"}

输入数据是个字典，包含 “data" 和 "version" 两个键值。data 是个数组，里面的每一项对应一本书，以及这本书中的的「语境，问题，答案」字典列表。

对于每个「语境，问题，答案」，其格式又如下：

{context and qas item 1} =
{"context": "正闹着，贾母遣人来叫他吃饭，方往前边来，胡乱吃了半碗，仍回自己房中．只见袭人睡在外头炕上，麝月在旁边抹骨牌．宝玉素知麝月与袭人亲厚，一并连麝月也不理，揭起软帘自往里间来．麝月只得跟进来．平儿便推他出去，说：",
"qas" : [ {"answers":[{"answer_start": 46, "text":"平儿"}],
"question": "接下来一句话是谁说的",
"id": "index"},
{question answer pair 2},
..., {question answer pair n}]
}

在这次尝试中，我只使用了经过 Data Augmentation 生成的200多万组数据中的 36000 组做训练。BERT 的 SQUAD 训练脚本 test_squad.sh 设置基本没改变，最大的改变是 max_seq_length=128，以及训练数据测试数据文件所在位置及内容，

export BERT_BASE_DIR="pathto/chinese_L-12_H-768_A-12"
export SQUAD_DIR="pathto/squad_data_chinese"
python pathto/run_squad.py \
--vocab_file=$BERT_BASE_DIR/vocab.txt \
--bert_config_file=$BERT_BASE_DIR/bert_config.json \
--init_checkpoint=$BERT_BASE_DIR/bert_model.ckpt \
--do_train=True \
--train_file=$SQUAD_DIR/chinese_speaker_squad.json \
--do_predict=True \
--predict_file=$SQUAD_DIR/chinese_speaker_squad_valid.json \
--train_batch_size=12 \
--learning_rate=3e-5 \
--num_train_epochs=2.0 \
--max_seq_length=128 \
--doc_stride=128 \
--output_dir=pathto/squad_data_chinese

（四）预测结果

因为BERT在维基百科的大量中文语料上做过训练，已经掌握了中文的基本规律。而少量的训练数据微调，即可让BERT知道它所需要处理的任务类型。通过简单的阅读理解与问答训练，说话人提取的任务效果惊人，虽然还没有人工完全验证提取结果的正确性，但是从语境和答案对看来，大部分结果无差错。总共数据是10683条，打了标签的训练数据是前面的1500多条。下面将预测的10683条中从后往前数的部分预测结果列出，

想了一回，也觉解了好些．又想到袭人身上：||| 袭人 (此预测结果❌）
那日薛姨妈并未回家，因恐宝钗痛哭，所以在宝钗房中解劝．那宝钗却是极明理，思前想后，宝玉原是一种奇异的人．夙世前因，自有一定，原无可怨天尤人．了．薛姨妈心里反倒安了，便到王夫人那里先把宝钗的话说了．王夫人点头叹道：||| 王夫人
说着，更又伤心起来．薛姨妈倒又劝了一会子，因又提起袭人来，说：||| 薛姨妈
王夫人道：||| 王夫人
薛姨妈道：||| 薛姨妈
王夫人听了道：||| 王夫人
薛姨妈听了点头道：||| 薛姨妈
看见袭人泪痕满面，薛姨妈便劝解譬喻了一会．W袭人本来老实，不是伶牙利齿的人，薛姨妈说一句，他应一句，回来说道：||| 薛姨妈（此结果从语境看不出是否正确）
过了几日，贾政回家，众人迎接．贾政见贾赦贾珍已都回家，弟兄叔侄相见，大家历叙别来的景况．然后内眷们见了，不免想起宝玉来，又大家伤了一会子心．贾政喝住道：||| 贾政
次日贾政进内，请示大臣们，说是：||| 贾政
回到家中，贾琏贾珍接着，贾政将朝内的话述了一遍，众人喜欢．贾珍便回说：||| 贾珍
贾政并不言语，隔了半日，却吩咐了一番仰报天恩的话．贾琏也趁便回说：||| 贾琏
贾政昨晚也知巧姐的始末，便说：||| 贾政
贾琏答应了"是"，又说：||| 贾琏
贾政道：||| 贾政
贾政说毕进内．贾琏打发请了刘姥姥来，应了这件事．刘姥姥见了王夫人等，便说些将来怎样升官，怎样起家，怎样子孙昌盛．正说着，丫头回道：||| 丫头
王夫人问几句话，花自芳的女人将亲戚作媒，说的是城南蒋家的，现在有房有地，又有铺面，姑爷年纪略大了几岁，并没有娶过的，况且人物儿长的是百里挑一的．王夫人听了愿意，说道：||| 王夫人
王夫人又命人打听，都说是好．王夫人便告诉了宝钗，仍请了薛姨妈细细的告诉了袭人．袭人悲伤不已，又不敢违命的，心里想起宝玉那年到他家去，回来说的死也不回去的话，"如今太太硬作主张．若说我守着，又叫人说我不害臊，若是去了，实不是我的心愿"，便哭得咽哽难鸣，又被薛姨妈宝钗等苦劝，回过念头想道：||| 薛姨妈宝钗(此预测结果❌）
于是，袭人含悲叩辞了众人，那姐妹分手时自然更有一番不忍说．袭人怀着必死的心肠上车回去，见了哥哥嫂子，也是哭泣，但只说不出来．那花自芳悉把蒋家的娉礼送给他看，又把自己所办妆奁一一指给他瞧，说那是太太赏的，那是置办的．袭人此时更难开口，住了两天，细想起来：||| 袭人
不言袭人从此又是一番天地．且说那贾雨村犯了婪索的案件，审明定罪，今遇大赦，褫籍为民．雨村因叫家眷先行，自己带了一个小厮，一车行李，来到急流津觉迷渡口．只见一个道者从那渡头草棚里出来，执手相迎．雨村认得是甄士隐，也连忙打恭，士隐道：||| 士隐
雨村道：||| 雨村
甄士隐道：||| 甄士隐
雨村欣然领命，两人携手而行，小厮驱车随后，到了一座茅庵．士隐让进雨村坐下，小童献上茶来．雨村便请教仙长超尘的始末．士隐笑道：||| 士隐
雨村道：||| 雨村
士隐道：||| 士隐
雨村惊讶道：||| 雨村
士隐道：||| 士隐
雨村道：||| 雨村
士隐道：||| 士隐
雨村听了，虽不能全然明白，却也十知四五，便点头叹道：||| 雨村
士隐笑道：||| 士隐
雨村听着，却不明白了．知仙机也不便更问，因又说道：||| 雨村听着，却不明白了．知仙机 (此预测结果❌）
士隐叹息道：||| 士隐
雨村听到这里，不觉拈须长叹，因又问道：||| 雨村
士隐道：||| 士隐
雨村低了半日头，忽然笑道：||| 雨村
士隐微微笑道：||| 士隐
食毕，雨村还要问自己的终身，士隐便道：||| 士隐
雨村惊讶道：||| 雨村
士隐道：||| 士隐
这士隐自去度脱了香菱，送到太虚幻境，交那警幻仙子对册，刚过牌坊，见那一僧一道，缥渺而来．士隐接着说道：||| 士隐
那僧说：||| 那僧
这一日空空道人又从青埂峰前经过，见那补天未用之石仍在那里，上面字迹依然如旧，又从头的细细看了一遍，见后面偈文后又历叙了多少收缘结果的话头，便点头叹道：||| 空空道人
想毕，便又抄了，仍袖至那繁华昌盛的地方，遍寻了一番，不是建功立业之人，即系饶口谋衣之辈，那有闲情更去和石头饶舌．直寻到急流津觉迷度口，草庵中睡着一个人，因想他必是闲人，便要将这抄录的《石头记》给他看看．那知那人再叫不醒．空空道人复又使劲拉他，才慢慢的开眼坐起，便草草一看，仍旧掷下道：||| 空空道人
空空道人忙问何人，那人道：||| 那人
那空空道人牢牢记着此言，又不知过了几世几劫，果然有个悼红轩，见那曹雪芹先生正在那里翻阅历来的古史．空空道人便将贾雨村言了，方把这《石头记》示看．那雪芹先生笑道：||| 雪芹先生
空空道人便问：||| 空空道人
曹雪芹先生笑道：||| 曹雪芹先生
那空空道人听了，仰天大笑，掷下抄本，飘然而去．一面走着，口中说道：||| 空空道人

结果分析：大部分简单的语境，BERT都可以正确的预测谁是说话的那个人，但是有些复杂一点的，就会出错，比如上面这些例子中的：

想了一回，也觉解了好些．又想到袭人身上：||| 袭人 (此预测结果❌）
王夫人又命人打听，都说是好．王夫人便告诉了宝钗，仍请了薛姨妈细细的告诉了袭人．袭人悲伤不已，又不敢违命的，心里想起宝玉那年到他家去，回来说的死也不回去的话，"如今太太硬作主张．若说我守着，又叫人说我不害臊，若是去了，实不是我的心愿"，便哭得咽哽难鸣，又被薛姨妈宝钗等苦劝，回过念头想道：||| 薛姨妈宝钗(此预测结果❌）
雨村听着，却不明白了．知仙机也不便更问，因又说道：||| 雨村听着，却不明白了．知仙机(此预测结果❌）

第三个错误最是搞笑，好像机器还没有明白“雨村听着，却不明白了．知仙机“并不是一个人的名字。

下面我再从其他预言的结果中挑选了一些看起来不容易预测，但是机器正确理解并预测的例子：

10575 贾兰那里肯走．尤氏等苦劝不止．众人中只有惜春心里却明白了，只不好说出来，便问宝钗道：||| 惜春
10183 王夫人已到宝钗那里，见宝玉神魂失所，心下着忙，便说袭人道：||| 王夫人
王仁便叫了他外甥女儿巧姐过来说：||| 王仁（下面一句话算谁说的？？？我也很懵）
9490 正推让着，宝玉也来请薛姨妈李婶娘的安．听见宝钗自己推让，他心里本早打算过宝钗生日，因家中闹得七颠八倒，也不敢在贾母处提起，今见湘云等众人要拜寿，便喜欢道：||| 宝玉

（四）人物关系分析

按照相邻的两个说话者极有可能是对话者统计出红楼梦中人物关系如下，宝玉与袭人之间对话最多（178+175），宝玉与黛玉之间对话次之（177+174），宝玉与宝钗之间对话（65+61），仅从对话次数来看，袭人与黛玉在宝玉心目中的占地差不多，宝钗（65+61）占地只相当于黛玉的三分之一，略高于晴雯（46+41）。

通过这个例子，深深感觉google 的BERT预训练+微调的自然语言处理模型之强大。很多nlp的问题可以转换成 “阅读理解 + 问答”（SQUAD）的问题。在此写下假期3天做的一个有趣的尝试，希望在知乎上看到知友们使用BERT开发出更多好玩的应用：）

[('宝玉-袭人', 178),
('黛玉-宝玉', 177),
('袭人-宝玉', 175),
('宝玉-黛玉', 174),
('宝玉-宝玉', 137),
('贾母-贾母', 115),
('宝玉-宝钗', 65),
('凤姐-凤姐', 64),
('宝钗-宝玉', 61),
('黛玉-黛玉', 59),
('贾母-凤姐', 57),
('贾政-贾政', 54),
('袭人-袭人', 48),
('宝玉-晴雯', 46),
('贾琏-凤姐', 46),
('宝钗-黛玉', 45),
('凤姐-贾母', 44),
('黛玉-宝钗', 42),
('凤姐-贾琏', 42),
('王夫人-贾母', 41),
('宝玉-贾母', 41),
('晴雯-宝玉', 41),
('王夫人-宝玉', 41),
('贾母-宝玉', 40),
('宝玉-贾政', 39),
('黛玉-紫鹃', 39),
('黛玉-湘云', 38),
('紫鹃-黛玉', 37),
('凤姐儿-贾母', 35),
('众人-贾政', 35)]

（五）条件随机场 CRF 的预言结果

除了这种暴力的 BERT阅读理解与问答，条件随机场CRF也可以用来做 Pos-Tagging，作为对比。CRF 的方法相比 BERT 微调的方法所耗资源极小，单机上几分钟就可以训练出来。从结果来看，大部分比较短的语境下CRF都可以很好的识别说话人。对于很长的语境，CRF会挑选出多个人物，而不仅仅是说话的人。比如下面这种，

见面时彼此悲喜交接，未免又大哭一阵，后又致喜庆之词．宝玉心中品度黛玉，越发出落的超逸了．黛玉又带了许多书籍来，忙着打扫卧室，安插器具，又将些纸笔等物分送宝钗，迎春，宝玉等人．宝玉又将北静王所赠 й к 香串珍重取出来，转赠黛玉．门子说：

CRF 同时挑出了 “黛玉”， “宝玉” 和 “门子” 三个选项。如果巧妙的设置一些特征函数，应该能解决这个问题吧。

CRF 的例子放在：conditional_random_field.ipynb 中，运行之后的输出结果如下图：

对比 BERT 与条件随机场（CRF）的预测结果，当语境中没有说话的人，CRF 预测的更好。比如下面这些预测结果：

BERT：因问：||| 因
CRF：因问：|||
BERT：又忙携黛玉之手，问：||| 黛玉
CRF：又忙携黛玉之手，问：|||
BERT：一面又问婆子们：||| 一面
CRF：一面又问婆子们：|||

但是当语境中有多个人物，CRF 会挑出很多个候选，并不像BERT一样，每条语境只挑出一个候选。

BERT：黛玉便说了名．宝玉又问表字．黛玉道：||| 黛玉
CRF：黛玉便说了名．宝玉又问表字．黛玉道：||| 黛玉宝玉黛玉
BERT：这里凤姐叫人抓些果子与板儿吃，刚问些闲话时，就有家下许多媳妇管事的来回话．平儿回了，凤姐道：||| 凤姐
CRF：这里凤姐叫人抓些果子与板儿吃，刚问些闲话时，就有家下许多媳妇管事的来回话．平儿回了，凤姐道：||| 凤姐平儿凤姐

而且有很多语境，BERT 能成功识别说话的人，但 CRF 莫名其妙的失败，可能是因为训练语料中没有出现这种模式。因为BERT在大规模语料上预训练过，对语言理解的更为深刻，所以泛化能力也更好，比如下面这些例子，

BERT：说话时，已摆了茶果上来．熙凤亲为捧茶捧果．又见二舅母问他：||| 二舅母
CRF：说话时，已摆了茶果上来．熙凤亲为捧茶捧果．又见二舅母问他：|||
BERT：当下茶果已撤，贾母命两个老嬷嬷带了黛玉去见两个母舅．时贾赦之妻邢氏忙亦起身，笑回道：||| 邢氏
CRF：当下茶果已撤，贾母命两个老嬷嬷带了黛玉去见两个母舅．时贾赦之妻邢氏忙亦起身，笑回道：|||

数了一下，总共10683条里面，BERT 与 CRF 预测结果不同的共有4818个，人工对比了预测结果不同的列表中前354个，其中BERT预言正确的是277个，CRF预言正确的是68个。预言结果相同的5865条基本全对。所以粗略的估计，BERT的说话人提取正确率为 (5865 + 4818x277/354.)/10683 约等于 90%, 而 CRF 的正确率为，（5865 + 4818x68/354.)/10683 约等于 64%。CRF 的结果中，有相当一部分错误原因是预测了多个说话人，目前尚不知如何约束CRF只输出说出下一句话的那个人。如果有 CRF 的高手，可以下载训练数据，看看能不能提高其准确率。

（六）小说主要角色词云

在网上看到很多制作漂亮的词云，自己试了下，如果直接用 jieba 分词，统计词频，制作出的词云图片里有很多不相关的词。而 BERT 既然帮我们识别出了所有说话的人，那么我们可以制作一个小说主要角色词云，只需要几行代码，

from PIL import Image
from wordcloud import WordCloud, ImageColorGenerator
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
sns.set_context('poster')def BertWordCloud(bert_speaker_path, img_path, save_as="wordcloud.png"):with open(bert_speaker_path, "r") as fin:lines = fin.readlines()segs = [line.split('|||')[1].strip() for line in lines]cut_text = ' '.join(segs)background_img = np.array(Image.open(img_path))wordcloud = WordCloud(font_path="/System/Library/fonts/PingFang.ttc",background_color="white",mask=background_img,collocations=False).generate(cut_text)# collocations=False will remove repeated keywordsimg_colors = ImageColorGenerator(background_img)plt.imshow(wordcloud.recolor(color_func=img_colors), interpolation="bilinear")plt.axis("off")plt.savefig(save_as)plt.show()BertWordCloud("res.txt", "alice.png", save_as="honglou_speaker.png")

结果如下，宝玉果然是当仁不让的主角，虽然袭人与宝玉之间的交流多于黛玉与宝玉的交流，但因为黛玉与其他角色之间有更多的交互，所以从词云上看，黛玉更接近主角。