Home

Awesome

NJUNMT-pytorch


English, 中文

License: MIT Build Status

NJUNMT-pytorch is an open-source toolkit for neural machine translation. This toolkit is highly research-oriented, which contains some common baseline model:

Table of Contents

Requirements

Usage

0. Quick Start

We provide push-button scripts to setup training and inference of transformer model on NIST Chinese-English Corpus (only on NJUNLP's server). Just execute under root directory of this repo

bash ./scripts/train.sh

for training and

# 3 means decoding on NIST 2003. This value
# can also be 4,5,6, which represents NIST 2004, 2005, 2006 respectively. 
bash ./scripts/translate.sh 3 

1. Build Vocabulary

First we should generate vocabulary files for both source and target language. We provide a script in ./data/build_dictionary.py to build them in json format.

See how to use this script by running:

python ./scripts/build_dictionary.py --help

We highly recommend not to set the limitation of the number of words and control it by config files while training.

2. Write Configuration File

See examples in ./configs folder. We provide several examples:

To further learn how to configure a NMT training task, see this wiki page.

3. Training

We can setup a training task by running

export CUDA_VISIBLE_DEVICES=0
python -m src.bin.train \
    --model_name <your-model-name> \
    --reload \
    --config_path <your-config-path> \
    --log_path <your-log-path> \
    --saveto <path-to-save-checkpoints> \
    --valid_path <path-to-save-validation-translation> \
    --use_gpu

See detail options by running python -m src.bin.train --help.

During training, checkpoints and best models will be saved under the directory specified by option ---saveto. Suppose that the model name is "MyModel", there would be several files under that directory:

4. Translation

When training is over, our code will automatically save the best model. Usually you could just use the final best model, which is named as xxxx.best.final, to translate. This model achieves the best performance on the validation set.

We can translation any text by running:

export CUDA_VISIBLE_DEVICES=0
python -m src.bin.translate \
    --model_name <your-model-name> \
    --source_path <path-to-source-text> \
    --model_path <path-to-model> \
    --config_path <path-to-configuration> \
    --batch_size <your-batch-size> \
    --beam_size <your-beam-size> \
    --alpha <your-length-penalty> \
    --use_gpu

See detail options by running python -m src.bin.translate --help.

Also our code support ensemble decoding. See more options by running python -m src.bin.ensemble_translate --help

Benchmark

See BENCHMARK.md

Contact

If you have any question, please contact whr94621@foxmail.com