🦄 使用迁移学习构建最先进的对话式AI

本仓库包含与博客文章 🦄 如何使用迁移学习构建最先进的对话式AI 相关的代码。

该代码库整洁且带有注释，包含训练和测试脚本，可用于基于 OpenAI GPT 和 GPT-2 Transformer 语言模型的迁移学习来训练对话代理。

此代码库可用于复现 HuggingFace 参加 NeurIPS 2018 对话竞赛 ConvAI2 的结果，该参赛作品在自动评估指标上处于当时最先进水平。超过 3000 行的竞赛代码被精简为约 250 行的训练代码，并支持分布式训练和 FP16 精度选项，最终形成了本仓库。

该模型可以在配备 8 块 V100 显卡的云实例上大约一小时内完成训练（当前成本约为 25 美元），同时我们也提供了预训练好的模型。

安装

要安装并使用训练和推理脚本，请克隆仓库并安装依赖项：

git clone https://github.com/huggingface/transfer-learning-conv-ai
cd transfer-learning-conv-ai
pip install -r requirements.txt
python -m spacy download en

使用 Docker 安装

若使用 Docker 进行安装，请构建自包含镜像：

docker build -t convai .

注意：请确保您的 Docker 设置为容器构建分配了足够的内存。使用默认的 1.75GB 内存进行构建会因 PyTorch 轮子文件过大而失败。

随后您可以进入镜像：

ip-192-168-22-157:transfer-learning-conv-ai loretoparisi$ docker run --rm -it convai bash
root@91e241bb823e:/# ls
Dockerfile  README.md  boot                  dev  home         lib    media  models  proc              root  sbin  sys  train.py  utils.py
LICENCE     bin        convai_evaluation.py  etc  interact.py  lib64  mnt    opt     requirements.txt  run   srv   tmp  usr       var

然后您可以运行 interact.py 脚本来使用预训练模型：

python3 interact.py --model models/

预训练模型

我们在 S3 上提供了一个预训练并微调过的模型，地址为 这里。下载并使用该模型最简单的方式就是直接运行 interact.py 脚本与模型对话。不带任何参数时，该脚本会自动下载并缓存我们的模型。

使用训练脚本

训练脚本可在单 GPU 或多 GPU 环境中运行：

python ./train.py  # 单 GPU 训练
python -m torch.distributed.launch --nproc_per_node=8 ./train.py  # 在 8 张 GPU 上训练

训练脚本接受多个参数以调整训练过程：

参数	类型	默认值	描述
dataset_path	str	""	数据集路径或 URL。若为空，则从 S3 下载。
dataset_cache	str	'./dataset_cache.bin'	数据集缓存路径或 URL。
model	str	"openai-gpt"	模型路径、URL 或简称。
num_candidates	int	2	训练时的候选回复数量。
max_history	int	2	保留的历史对话轮数。
train_batch_size	int	4	训练批次大小。
valid_batch_size	int	4	验证批次大小。
gradient_accumulation_steps	int	8	多步梯度累积。
lr	float	6.25e-5	学习率。
lm_coef	float	1.0	语言模型损失系数。
mc_coef	float	1.0	多选题损失系数。
max_norm	float	1.0	梯度裁剪阈值。
n_epochs	int	3	训练轮数。
personality_permutations	int	1	个性句子的排列组合数。
device	str	"cuda" if torch.cuda.is_available() else "cpu"	设备（CUDA 或 CPU）。
fp16	str	""	设置为 O0、O1、O2 或 O3 以启用 FP16 训练（参见 Apex 文档）。
local_rank	int	-1	分布式训练中的本地排名（-1：非分布式）。

以下是在拥有 8 块 V100 GPU 的服务器上复现我们结果的方法（请根据您的配置调整节点数和批次大小）：

python -m torch.distributed.launch --nproc_per_node=8 ./train.py --gradient_accumulation_steps=4 --lm_coef=2.0 --max_history=2 --n_epochs=1 --num_candidates=4 --personality_permutations=2 --train_batch_size=2 --valid_batch_size=2

使用 convai2 评估脚本，该模型应能获得 Hits@1 超过 79、困惑度 20.5 和 F1 16.5 的成绩（见下文）。

这些数值略低于我们在 ConvAI2 竞赛中取得的成绩。若想达到相同效果，可以尝试以下调整：

• 在 ConvAI2 竞赛中，我们还对位置嵌入进行了调整，使对话历史始终从相同的嵌入开始。这可以通过 pytorch-transformers 轻松实现，并有望提升 Hits@1 指标。
• 在 ConvAI2 竞赛中，我们使用了束搜索解码器。虽然在 F1 指标上表现更好，但我们认为，与本仓库提供的核采样解码器相比，束搜索带来的用户体验稍显不足。

使用交互脚本

训练脚本会将所有实验和检查点保存在仓库根目录下的 ./runs 文件夹中，并以实验时间戳命名子文件夹。

随后，您可以通过指向该文件夹来使用交互脚本与模型互动。

以下是运行交互脚本的示例命令：

python ./interact.py --model_checkpoint ./data/Apr17_13-31-38_thunder/  # 使用训练检查点运行交互脚本
python ./interact.py  # 使用我们 S3 上的微调模型运行交互脚本

微调后的模型将给出 FINAL Hits@1: 0.715。

交互脚本接受一些参数来调整解码算法：

参数	类型	默认值	描述
dataset_path	str	""	数据集路径或 URL。若为空，则从 S3 下载。
dataset_cache	str	'./dataset_cache.bin'	数据集缓存路径或 URL。
model	str	"openai-gpt"	模型路径、URL 或简称。
max_history	int	2	保留的历史对话轮数。
device	str	cuda 如果 torch.cuda.is_available()，否则为 cpu	设备（CUDA 或 CPU）。
no_sample	action store_true	设置为真以使用贪婪解码而非采样。
max_length	int	20	输出回复的最大长度。
min_length	int	1	输出回复的最小长度。
seed	int	42	随机种子。
temperature	int	0.7	采样 softmax 温度。
top_k	int	0	采样前过滤 top-k 个 token（<=0：不进行过滤）。
top_p	float	0.9	采样前进行核过滤（top-p）（<=0.0：不进行过滤）。

运行 ConvAI2 评估脚本

要运行 ConvAI2 挑战赛的评估脚本，您首先需要在仓库的根目录下安装 ParlAI，方法如下：

git clone https://github.com/facebookresearch/ParlAI.git
cd ParlAI
python setup.py develop

然后，您可以在 ParlAI 的根目录下运行评估脚本：

cd ParlAI
python ../convai_evaluation.py --eval_type hits@1  # 下载并使用 hits@1 指标评估我们微调后的模型
python ../convai_evaluation.py --eval_type hits@1  --model_checkpoint ./data/Apr17_13-31-38_thunder/  # 使用 hits@1 指标评估某个训练检查点

评估脚本接受几个参数来选择评估指标并调整解码算法：

参数	类型	默认值	描述
eval_type	str	"hits@1"	在 ConvAI2 验证集上使用 hits@1、ppl 或 f1 指标评估模型
model	str	"openai-gpt"	模型的路径、URL 或简称
max_history	int	2	保留对话历史中的前几轮发言数量
device	str	cuda（如果 torch.cuda.is_available()）否则为 cpu	设备（cuda 或 cpu）
no_sample	store_true	设置为使用贪婪解码而非采样
max_length	int	20	输出发言的最大长度
min_length	int	1	输出发言的最小长度
seed	int	42	随机种子
temperature	float	0.7	采样 softmax 的温度
top_k	int	0	采样前过滤 top-k 个词（<=0：不进行过滤）
top_p	float	0.9	采样前进行核采样过滤（<=0.0：不进行过滤）

数据格式

请参阅 example_entry.py 文件及其顶部的注释。

引用

如果您在研究中使用了这段代码，可以引用我们在 NeurIPS CAI 工作组会议上的论文 paper：

@article{DBLP:journals/corr/abs-1901-08149,
  author    = {Thomas Wolf and
               Victor Sanh and
               Julien Chaumond and
               Clement Delangue},
  title     = {TransferTransfo: {A} Transfer Learning Approach for Neural Network
               Based Conversational Agents},
  journal   = {CoRR},
  volume    = {abs/1901.08149},
  year      = {2019},
  url       = {http://arxiv.org/abs/1901.08149},
  archivePrefix = {arXiv},
  eprint    = {1901.08149},
  timestamp = {Sat, 02 Feb 2019 16:56:00 +0100},
  biburl    = {https://dblp.org/rec/bib/journals/corr/abs-1901-08149},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}

transfer-learning-conv-ai 快速上手指南

本指南帮助开发者快速部署基于迁移学习的对话 AI 模型（OpenAI GPT/GPT-2），复现 HuggingFace 在 ConvAI2 竞赛中的成果。

环境准备

• 操作系统: Linux / macOS (Windows 需使用 Docker 或 WSL)
• 硬件要求:
  • 训练: 推荐多卡环境（如 8x V100），单卡亦可但耗时较长。
  • 推理/交互: 单张 GPU 或 CPU 即可运行预训练模型。
• 软件依赖:
  • Python 3.6+
  • PyTorch
  • Git
  • Docker (可选，用于容器化部署)

国内加速建议： 若下载依赖或模型较慢，可配置以下环境变量使用国内镜像：

export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
export HF_ENDPOINT=https://hf-mirror.com

安装步骤

方式一：本地安装（推荐）

1. 克隆仓库

   git clone https://github.com/huggingface/transfer-learning-conv-ai
   cd transfer-learning-conv-ai
   

2. 安装依赖

   pip install -r requirements.txt
   python -m spacy download en
   

方式二：Docker 安装

1. 构建镜像

   注意：请确保 Docker 分配给容器的内存至少为 4GB（默认 1.75GB 会导致构建失败）。

   docker build -t convai .
   

2. 启动容器

   docker run --rm -it convai bash
   

基本使用

1. 与预训练模型交互（最简单用法）

无需手动下载模型，脚本会自动从 S3 下载并缓存微调后的模型。

本地运行：

python interact.py

Docker 内运行：

python3 interact.py --model models/

运行后直接在终端输入即可与机器人对话。

2. 训练自定义模型

单卡训练：

python ./train.py

多卡分布式训练（例如 8 卡）：

python -m torch.distributed.launch --nproc_per_node=8 ./train.py

复现竞赛高性能结果（8x V100 配置参考）：

python -m torch.distributed.launch --nproc_per_node=8 ./train.py --gradient_accumulation_steps=4 --lm_coef=2.0 --max_history=2 --n_epochs=1 --num_candidates=4 --personality_permutations=2 --train_batch_size=2 --valid_batch_size=2

3. 使用训练好的检查点进行交互

训练脚本会将模型保存在 ./runs 目录下（以时间戳命名文件夹）。

# 替换为你的实际检查点路径
python ./interact.py --model_checkpoint ./data/Apr17_13-31-38_thunder/

4. 模型评估 (ConvAI2 指标)

如需计算 Hits@1, Perplexity 或 F1 分数，需先安装 ParlAI：

git clone https://github.com/facebookresearch/ParlAI.git
cd ParlAI
python setup.py develop

返回项目根目录运行评估：

cd ..
# 评估预训练模型
python convai_evaluation.py --eval_type hits@1

# 评估本地训练的检查点
python convai_evaluation.py --eval_type hits@1 --model_checkpoint ./data/Apr17_13-31-38_thunder/