bandit-nmt

THIS REPO DEMONSTRATES HOW TO INTEGRATE A POLICY GRADIENT METHOD INTO NMT. FOR A STATE-OF-THE-ART NMT CODEBASE, VISIT simple-nmt.

This is code repo for our EMNLP 2017 paper "Reinforcement Learning for Bandit Neural Machine Translation with Simulated Human Feedback", which implements the A2C algorithm on top of a neural encoder-decoder model and benchmarks the combination under simulated noisy rewards.

Requirements: - Python 3.6 - PyTorch 0.2

NOTE: as of Sep 16 2017, the code got 2x slower when I upgraded to PyTorch 2.0. This is a known issue and PyTorch is fixing it.

IMPORTANT: Set home directory (otherwise scripts will not run correctly): ~~~~

export BANDIT_HOME=$PWD export DATA=$BANDIT_HOME/data export SCRIPT=$BANDIT_HOME/scripts ~~~~

Data extraction

Download pre-processing scripts ~~~~

cd $DATA/scripts bash download_scripts.sh ~~~~

For German-English ~~~~

cd $DATA/en-de bash extract_data_de_en.sh ~~~~

NOTE: train_2014 and train_2015 highly overlap. Please be cautious when using them for other projects.

Data should be ready in $DATA/en-de/prep

TODO: Chinese-English needs segmentation

Data pre-processing

~~~~

cd $SCRIPT bash make_data.sh de en ~~~~

Pretraining

Pretrain both actor and critic ~~~~

cd $SCRIPT bash pretrain.sh en-de $YOUR_LOG_DIR ~~~~

See scripts/pretrain.sh for more details.

Pretrain actor only ~~~~

cd $BANDIT_HOME python train.py -data $YOUR_DATA -save_dir $YOUR_SAVE_DIR -end_epoch 10 ~~~~

Reinforcement training

~~~~

cd $BANDIT_HOME ~~~~

From scratch ~~~~

python train.py -data $YOUR_DATA -save_dir $YOUR_SAVE_DIR -start_reinforce 10 -end_epoch 100 -critic_pretrain_epochs 5 ~~~~

From a pretrained model ~~~~

python train.py -data $YOUR_DATA -load_from $YOUR_MODEL -save_dir $YOUR_SAVE_DIR -start_reinforce -1 -end_epoch 100 -critic_pretrain_epochs 5 ~~~~

Perturbed rewards

For example, use thumb up/thump down reward: ~~~~

cd $BANDIT_HOME python train.py -data $YOUR_DATA -load_from $YOUR_MODEL -save_dir $YOUR_SAVE_DIR -start_reinforce -1 -end_epoch 100 -critic_pretrain_epochs 5 -pert_func bin -pert_param 1 ~~~~

See lib/metric/PertFunction.py for more types of function.

Evaluation

~~~~

cd $BANDIT_HOME ~~~~

On heldout sets (heldout BLEU): ~~~~

python train.py -data $YOUR_DATA -load_from $YOUR_MODEL -eval -save_dir . ~~~~

On bandit set (per-sentence BLEU): ~~~~

python train.py -data $YOUR_DATA -load_from $YOUR_MODEL -eval_sample -save_dir . ~~~~