在上一期教程中，我演示了如何从零开始训练一个比较智障的聊天机器人。
但是当时数据量太少，模型简单，完全没法用，只能回复训练集中出现过的句子。
而现在，完全体的聊天机器人终于诞生了，我把它命名为“杨超越”。这次可以随你提问什么了，她都能对答如流！
首先来看看回复的效果怎么样：
聊天开始！（按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.model和spm.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下，里面还细分为了best、best_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加速库都非常好用，从上面使用教程中也可以看出，几乎不需要你写什么代码就能使用起来。
