养成女友?我训练出了一个“杨超越”聊天机器人

在上一期教程中,我演示了如何从零开始训练一个比较智障的聊天机器人。

https://zhuanlan.zhihu.com/p/392175369

但是当时数据量太少,模型简单,完全没法用,只能回复训练集中出现过的句子。

而现在,完全体的聊天机器人终于诞生了,我把它命名为“杨超越”。这次可以随你提问什么了,她都能对答如流!

首先来看看回复的效果怎么样:

聊天开始!(按q退出)
我:很高兴认识你
杨超越:我也很开心哦
我:我喜欢你
杨超越:我也喜欢你
我:做我女朋友好不好?
杨超越:哈哈,可以呢!
我:我会一直支持超越妹妹的!
杨超越:谢谢我们会一直努力的
我:我什么时候有女朋友?
杨超越:女朋友。。
我:我什么时候脱单?
杨超越:脱你妹啊!!
我:q
聊天结束!

可以看到超越妹妹的回复还是非常流畅的,那她究竟是怎么诞生的呢?

完整的项目我都放在了下面,欢迎大家点个star,支持一波:
https://github.com/godweiyang/chatbot

介绍

这里我才用的是网上公开的小黄鸡聊天语料,大概有100万条左右,但是质量不是很高,都放在了data目录下。

模型采用标准的Transformer-big模型,输入你的提问句子,预测超越妹妹回复的句子,config目录下是训练和预测的配置文件。

模型训练采用NeurST训练库,主要基于TensorFlow,也支持PyTorch训练。模型快速推理采用LightSeq,可加速推理10倍以上,同时还能加速NeurST的训练,最高加速3倍。两者都是字节跳动AI Lab自研的,都已开源。

安装环境

我们需要安装三样东西:

  • SentencePiece的命令行版本和python版本,用来对句子进行分词。
  • NeurST深度学习训练库,用来训练Transformer模型。
  • LightSeq,用来加速模型推理。

安装命令都很简单:

git clone https://github.com/google/sentencepiece.git & cd sentencepiece
mkdir build & cd build
cmake ..
make -j $(nproc)
sudo make install
sudo ldconfig -v

pip3 install lightseq neurst sentencepiece

开始养成

生成词表

首先我们需要从训练语料库中抽取出词表,为了方便,直接用SentencePiece来分词,生成大小为32k的词表。

spm_train --input=./data/train/train.src,./data/train/train.trg \
    --model_prefix=./data/spm \
    --vocab_size=32000 \
    --character_coverage=0.9995

这里需要指定训练语料路径--input、词表保存的路径前缀--model_prefix和词表大小--vocab_size。运行结束后会在data目录下生成spm.modelspm.vocab两个词表文件。一个是训练好的分词模型,一个是词表。

不过我也上传了生成好的TFRecord,大家也可以直接使用,跳过这一步。

我上传了生成好的词表文件,大家可以直接使用,跳过这一步。

生成TFRecord

为了加快TensorFlow的训练速度,可以预先将训练语料用上面的词表处理成id,然后保存为TFRecord格式。这样模型训练时就可以直接读取id进行训练了,不需要做前面的分词操作。能大大加快训练速度,提升显卡利用率。

python3 -m neurst.cli.create_tfrecords \
    --config_paths configs/task_args.yml \
    --dataset ParallelTextDataset \
    --src_file ./data/train/train.src \
    --trg_file ./data/train/train.trg \
    --processor_id 0 \
    --num_processors 1 \
    --num_output_shards 32 \
    --output_range_begin 0 \
    --output_range_end 32 \
    --output_template ./data/tfrecords/train.tfrecords-%5.5d-of-%5.5d

这里主要需要指定训练集的路径--src_file--trg_file,其它参数保持默认即可。生成完毕后会在data/tfrecords下面生成32个二进制文件,这就是处理好的训练数据了。

我上传了生成好的TFRecord,大家可以直接使用,跳过这一步。

模型训练

有了词表,有了处理好的训练数据,接下来就是训练模型了。这里开启了XLA优化,使用Horovod分布式训练,加快训练速度。如果报错,可以去掉最后两行。

python3 -m neurst.cli.run_exp \
    --entry trainer \
    --task translation \
    --hparams_set transformer_big \
    --model_dir ./models/transformer_big \
    --config_paths ./configs/task_args.yml,./configs/train_args.yml,./configs/valid_args.yml \
    --distribution_strategy horovod \
    --enable_xla

这里需要指定的参数就是模型保存路径model_dir,其他都保持默认。训练好的模型会保存在models/transformer_big下,里面还细分为了bestbest_avg等文件夹,用来存最好的模型、模型的平均值等等。

我在8张V100 32G显卡上训练了8个小时左右,如果你们自己训练的话还是比较耗时的。

由于模型文件过大,之后我会找地方上传我训练好的模型文件,省去大家训练的时间。

模型预测

训练好的模型会保存在models/transformer_big目录下,然后我们就可以开始预测啦。

python3 -m neurst.cli.run_exp \
    --entry predict \
    --model_dir ./models/transformer_big \
    --config_paths ./configs/predict_args.yml \
    --output output.txt

但是这时候还没有交互功能,只能指定一个测试集文件,写在了模型预测的配置文件里configs/predict_args.yml。还可以指定--output,将回复结果输出到文件中。

如果想直接体验交互式的对话聊天,可以跳过这一步。

模型导出为PB格式

如果直接用TensorFlow进行推理的话,速度非常慢,你就会感觉你和超越妹妹之间存在延时。所以可以将训练得到的ckpt模型导出为PB格式,然后就可以用LightSeq训练加速引擎进行快速推理了。

python3 export/export.py \
    --model_dir ./models/transformer_big \
    --output_file ./models/transformer_big/model.pb \
    --beam_size 4 \
    --length_penalty 0.6

这里需要指定模型路径--model_dir和导出PB文件的路径--output_file,其它参数保持默认。最后会得到models/transformer_big/model.pb这个PB文件。

由于模型文件过大,之后我会找地方上传我导出好的PB模型文件,这样大家就可以直接跳到最后一步了。

开始交互式聊天!

有了PB模型文件,就可以和超越妹妹开始聊天啦!

python3 chat.py \
    --spm_model ./data/spm.model \
    --model_file ./models/transformer_big/model.pb

这里需要指定两个路径。一是最开始训练好的分词模型--spm_model,用来将你输入的句子变成整数id。二是--model_file,也就是上一步中的PB格式模型文件。

聊天过程中随时可以按q退出聊天,你每说一句话,超越妹妹就会回复你一句。

欢迎关注

这次用到的NeurST训练库和LightSeq加速库都非常好用,从上面使用教程中也可以看出,几乎不需要你写什么代码就能使用起来。

聊天机器人:
https://github.com/godweiyang/chatbot

NeurST训练库:
https://github.com/bytedance/neurst

LightSeq加速库:
https://github.com/bytedance/lightseq


   转载规则


《养成女友?我训练出了一个“杨超越”聊天机器人》 韦阳 采用 知识共享署名 4.0 国际许可协议 进行许可。
 上一篇
分析完吴亦凡的所有歌曲,我震惊了 分析完吴亦凡的所有歌曲,我震惊了
最近吴亦凡的事件闹的沸沸扬扬,相信不少吃瓜群众也都有所耳闻。 但咱们今天不谈他的事件到底怎么样,咱们来谈谈他的歌到底如何。 收集歌词我写了一个代码,从QQ音乐上收集了吴亦凡所有的歌词信息,自动汇总整理成一个txt文件,代码在下面:https
下一篇 
从零开始训练一个人工智障女友 从零开始训练一个人工智障女友
很多人工智能小白可能不知道那些高大上的语音助理、机器翻译或者聊天机器人都是怎么被创造出来的,也不知道一个深度学习模型是怎么从零开始搭建并运行起来的。 今天我就简单教大家如何从零开始搭建一个Transformer模型,并在自己的数据上训练起来
2021-07-21
  目录