Home

Awesome

🦄 Building a State-of-the-Art Conversational AI with Transfer Learning

The present repo contains the code accompanying the blog post 🦄 How to build a State-of-the-Art Conversational AI with Transfer Learning.

This code is a clean and commented code base with training and testing scripts that can be used to train a dialog agent leveraging transfer Learning from an OpenAI GPT and GPT-2 Transformer language model.

This codebase can be used to reproduce the results of HuggingFace's participation to NeurIPS 2018 dialog competition ConvAI2 which was state-of-the-art on the automatic metrics. The 3k+ lines of competition code was distilled in about 250 lines of training code with distributed & FP16 options to form the present repository.

This model can be trained in about one hour on a 8 V100 cloud instance (currently costs about $25) and a pre-trained model is also made available.

Installation

To install and use the training and inference scripts please clone the repo and install the requirements:

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

Installation with Docker

To install using docker please build the self-contained image:

docker build -t convai .

Note: Make sure your Docker setup allocates enough memory to building the container. Building with the default of 1.75GB will fail due to large Pytorch wheel.

You can then enter the image

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

You can then run the interact.py script on the pretrained model:

python3 interact.py --model models/

Pretrained model

We make a pretrained and fine-tuned model available on our S3 here. The easiest way to download and use this model is just to run the interact.py script to talk with the model. Without any argument, this script will automatically download and cache our model.

Using the training script

The training script can be used in single GPU or multi GPU settings:

python ./train.py  # Single GPU training
python -m torch.distributed.launch --nproc_per_node=8 ./train.py  # Training on 8 GPUs

The training script accept several arguments to tweak the training:

ArgumentTypeDefault valueDescription
dataset_pathstr""Path or url of the dataset. If empty download from S3.
dataset_cachestr'./dataset_cache.bin'Path or url of the dataset cache
modelstr"openai-gpt"Path, url or short name of the model
num_candidatesint2Number of candidates for training
max_historyint2Number of previous exchanges to keep in history
train_batch_sizeint4Batch size for training
valid_batch_sizeint4Batch size for validation
gradient_accumulation_stepsint8Accumulate gradients on several steps
lrfloat6.25e-5Learning rate
lm_coeffloat1.0LM loss coefficient
mc_coeffloat1.0Multiple-choice loss coefficient
max_normfloat1.0Clipping gradient norm
n_epochsint3Number of training epochs
personality_permutationsint1Number of permutations of personality sentences
devicestr"cuda" if torch.cuda.is_available() else "cpu"Device (cuda or cpu)
fp16str""Set to O0, O1, O2 or O3 for fp16 training (see apex documentation)
local_rankint-1Local rank for distributed training (-1: not distributed)

Here is how to reproduce our results on a server with 8 V100 GPUs (adapt number of nodes and batch sizes to your configuration):

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

This model should give a Hits@1 over 79, perplexity of 20.5 and F1 of 16.5 using the convai2 evaluation script (see below).

These numbers are slightly lower than the number we obtained in the ConvAI2 competition. Here is what you can tweak to reach the same results:

Using the interaction script

The training script saves all the experiments and checkpoints in a sub-folder named with the timestamp of the experiment in the ./runs folder of the repository base folder.

You can then use the interactive script to interact with the model simply by pointing to this folder.

Here is an example command line to run the interactive script:

python ./interact.py --model_checkpoint ./data/Apr17_13-31-38_thunder/  # run the interactive script with a training checkpoint
python ./interact.py  # run the interactive script with the finetuned model on our S3

The fine-tuned model will gives FINAL Hits@1: 0.715

The interactive script accept a few arguments to tweak the decoding algorithm:

ArgumentTypeDefault valueDescription
dataset_pathstr""Path or url of the dataset. If empty download from S3.
dataset_cachestr'./dataset_cache.bin'Path or url of the dataset cache
modelstr"openai-gpt"Path, url or short name of the model
max_historyint2Number of previous utterances to keep in history
devicestrcuda if torch.cuda.is_available() else cpuDevice (cuda or cpu)
no_sampleaction store_trueSet to use greedy decoding instead of sampling
max_lengthint20Maximum length of the output utterances
min_lengthint1Minimum length of the output utterances
seedint42Seed
temperatureint0.7Sampling softmax temperature
top_kint0Filter top-k tokens before sampling (<=0: no filtering)
top_pfloat0.9Nucleus filtering (top-p) before sampling (<=0.0: no filtering)

Running ConvAI2 evaluation scripts

To run the evaluation scripts of the ConvAI2 challenge, you first need to install ParlAI in the repo base folder like this:

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

You can then run the evaluation script from ParlAI base folder:

cd ParlAI
python ../convai_evaluation.py --eval_type hits@1  # to download and evaluate our fine-tuned model on hits@1 metric
python ../convai_evaluation.py --eval_type hits@1  --model_checkpoint ./data/Apr17_13-31-38_thunder/  # to evaluate a training checkpoint on hits@1 metric

The evaluation script accept a few arguments to select the evaluation metric and tweak the decoding algorithm:

ArgumentTypeDefault valueDescription
eval_typestr"hits@1"Evaluate the model on hits@1, ppl or f1 metric on the ConvAI2 validation dataset
modelstr"openai-gpt"Path, url or short name of the model
max_historyint2Number of previous utterances to keep in history
devicestrcuda if torch.cuda.is_available() else cpuDevice (cuda or cpu)
no_sampleaction store_trueSet to use greedy decoding instead of sampling
max_lengthint20Maximum length of the output utterances
min_lengthint1Minimum length of the output utterances
seedint42Seed
temperatureint0.7Sampling softmax temperature
top_kint0Filter top-k tokens before sampling (<=0: no filtering)
top_pfloat0.9Nucleus filtering (top-p) before sampling (<=0.0: no filtering)

Data Format

see example_entry.py, and the comment at the top.

Citation

If you use this code in your research, you can cite our NeurIPS CAI workshop 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}
}