之前看到過一篇文章，通過提取文章中對話的人物，分析人物之間的關係，很好奇如何通過程式設計的方式知道一句話是誰說的。但是遍搜網路沒有發現類似的研究。
前段時間看到一個微信裡的讀書小程式，將人物對話都提取出來，將一本書的內容通過微信對話的方式表達出來，通過將對話的主角替換成讀者的微訊號以及使用者頭像，從而增加讀者的代入感。試了之後非常佩服程式作者的巧思。這使得我寫一個自然語言處理程式，提取書中對話，以及對話人物的念頭更加強烈。
之前並沒有多少 NLP 的經驗，只零碎試過用 LSTM 訓練寫唐詩，用 jieba 做分詞，用 Google 的 gensim 在 WikiPedia 中文語料上訓練詞向量。最近 Google 的 BERT 模型很火，執行了 BERT 的 SQuAD 閱讀理解與問答系統，分類器以及特徵提取例子之後，覺得這個任務可以用 BERT 微調來完成，在這裡記錄實驗的粗略步驟，與君共勉。
我把訓練資料和準備資料的指令碼開源，放在 GitLab 上，開放下載。
該目錄包含以下內容：
用於提取對話人物語境的指令碼 conversation_extraction.ipynb；
輔助打標籤的指令碼 label_data_by_click_buttons.ipynb；
提取出的語境檔案：honglou.py；
打過標籤的訓練資料：label_honglou.txt；
從打過標籤的資料合成百萬級別新資料的指令碼：augment_data.py；
將訓練資料轉換為 BERT/SQUAD 可讀的指令碼：prepare_squad_data.py；
預測結果檔案：res.txt（使用 36000 組資料訓練後的預測結果）；
預測結果檔案：res_1p2million.txt（使用 120萬 組資料訓練後的預測結果）。
對比之後發現使用更多的資料訓練所提升的效果有限，比較大的提升是後者在沒有答案時，輸出是輸入的完整拷貝。
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)]