storytold/storyteller-ml
2
0
Fork 0
mirror of https://github.com/storytold/storyteller-ml.git synced 2026-10-23 08:26:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Introduce code from Rust monorepo. This is a mess

Browse Source
This commit is contained in:
Brandon Thomas
2022-02-28 05:43:54 -05:00
parent 921fb49f9a
commit ba326bfcfd
1058 changed files with 192192 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tts/NeMo/examples/nlp/machine_translation/conf/aayn_base.yaml
+156
View File
@@ -0,0 +1,156 @@
name: AttentionIsAllYouNeed
do_training: True # set to False if only preprocessing data
do_testing: False # set to True to run evaluation on test data after training
model:
beam_size: 4
len_pen: 0.6
multilingual: False
max_generation_delta: 5
label_smoothing: 0.1
shared_tokenizer: True # train tokenizer model across src and tgt train data
preproc_out_dir: null # path to store data preprocessing outputs
src_language: 'en'
tgt_language: 'de'
train_ds:
src_file_name: null
tgt_file_name: null
use_tarred_dataset: False # if true tar_file_name and meta_file_name will be used (or created automatically)
# config for preprocessing training data and creating a tarred datset automatically
tar_file_prefix: parallel # prefix for tar file names
tar_files: null # if data has already been preprocessed (rest of config ignored)
metadata_file: null # metadata for tarred dataset
lines_per_dataset_fragment: 1000000 # Number of lines to consider for bucketing and padding
num_batches_per_tarfile: 100 # Number of batches (pickle files) within each tarfile
tar_shuffle_n: 100 # How many samples to look ahead and load to be shuffled
shard_strategy: scatter # tarred dataset shard distribution strategy
n_preproc_jobs: -2 # number of processes to use for data preprocessing (-2 means all but 2)
tokens_in_batch: 512
clean: true
max_seq_length: 512
shuffle: true
num_samples: -1
drop_last: false
pin_memory: false
num_workers: 8
concat_sampling_technique: temperature # only used with ConcatTranslationDataset
concat_sampling_temperature: 5 # only used with ConcatTranslationDataset
concat_sampling_probabilities: null # only used with ConcatTranslationDataset
validation_ds:
src_file_name: ???
tgt_file_name: ???
tokens_in_batch: 512
clean: false
max_seq_length: 512
shuffle: false
num_samples: -1
drop_last: false
pin_memory: false
num_workers: 8
test_ds:
src_file_name: ???
tgt_file_name: ???
tokens_in_batch: 512
clean: false
max_seq_length: 512
shuffle: false
num_samples: -1
drop_last: false
pin_memory: false
num_workers: 8
optim:
name: adam
lr: 0.001
betas:
- 0.9
- 0.98
weight_decay: 0.0
sched:
name: InverseSquareRootAnnealing
min_lr: 0.0
last_epoch: -1
warmup_ratio: 0.1
encoder_tokenizer:
library: yttm
tokenizer_model: null
vocab_size: null # vocab size for training bpe
bpe_dropout: null
vocab_file: null
special_tokens: null
training_sample_size: null # valid for sentencepiece tokenizer
decoder_tokenizer:
library: yttm
tokenizer_model: null
vocab_size: null # vocab size for training bpe
bpe_dropout: null
vocab_file: null
special_tokens: null
training_sample_size: null # valid for sentencepiece tokenizer
encoder:
library: nemo
model_name: null
pretrained: false
max_sequence_length: 512
num_token_types: 0
embedding_dropout: 0.1
learn_positional_encodings: false
hidden_size: 512
num_layers: 6
inner_size: 2048
num_attention_heads: 8
ffn_dropout: 0.1
attn_score_dropout: 0.1
attn_layer_dropout: 0.1
hidden_act: relu
mask_future: false
pre_ln: false
pre_ln_final_layer_norm: true
decoder:
library: nemo
model_name: null
pretrained: false
max_sequence_length: 512
num_token_types: 0
embedding_dropout: 0.1
learn_positional_encodings: false
hidden_size: 512
inner_size: 2048
num_layers: 6
num_attention_heads: 8
ffn_dropout: 0.1
attn_score_dropout: 0.1
attn_layer_dropout: 0.1
hidden_act: relu
pre_ln: false
pre_ln_final_layer_norm: true
head:
num_layers: 1
activation: relu
log_softmax: true
dropout: 0.0
use_transformer_init: true
trainer:
gpus: 4
num_nodes: 1
max_epochs: 200
amp_level: O2 # O1/O2 for mixed precision
precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
accelerator: ddp
checkpoint_callback: False
logger: False
log_every_n_steps: 50 # Interval of logging.
check_val_every_n_epoch: 1
exp_manager:
name: AAYNBase
files_to_copy: []
tts/NeMo/examples/nlp/machine_translation/conf/huggingface.yaml
+132
View File
@@ -0,0 +1,132 @@
name: HuggingFaceEncoder
do_training: True # set to False if only preprocessing data
do_testing: False # set to True to run evaluation on test data after training
model:
beam_size: 4
len_pen: 0.6
max_generation_delta: 5
label_smoothing: 0.1
shared_tokenizer: false
preproc_out_dir: null
src_language: 'en'
tgt_language: 'de'
train_ds:
src_file_name: null
tgt_file_name: null
use_tarred_dataset: False # if true tar_file_name and meta_file_name will be used (or created automatically)
# config for preprocessing training data and creating a tarred datset automatically
tar_file_prefix: parallel # prefix for tar file names
tar_files: null # if data has already been preprocessed (rest of config ignored)
metadata_file: null # metadata for tarred dataset
lines_per_dataset_fragment: 1000000 # Number of lines to consider for bucketing and padding
num_batches_per_tarfile: 100 # Number of batches (pickle files) within each tarfile
tar_shuffle_n: 100 # How many samples to look ahead and load to be shuffled
shard_strategy: scatter # tarred dataset shard distribution strategy
n_preproc_jobs: -2 # number of processes to use for data preprocessing (-2 means all but 2)
tokens_in_batch: 512
clean: true
max_seq_length: 512
shuffle: true
num_samples: -1
drop_last: false
pin_memory: false
num_workers: 8
validation_ds:
src_file_name: null
tgt_file_name: null
tokens_in_batch: 512
clean: false
max_seq_length: 512
shuffle: false
num_samples: -1
drop_last: false
pin_memory: false
num_workers: 8
test_ds:
src_file_name: null
tgt_file_name: null
tokens_in_batch: 512
clean: false
max_seq_length: 512
shuffle: false
num_samples: -1
drop_last: false
pin_memory: false
num_workers: 8
optim:
name: adam
lr: 0.001
betas:
- 0.9
- 0.98
weight_decay: 0.0
sched:
name: InverseSquareRootAnnealing
min_lr: 0.0
last_epoch: -1
warmup_ratio: 0.1
encoder_tokenizer:
library: huggingface
tokenizer_model: null
vocab_file: null
special_tokens: null
vocab_size: null
decoder_tokenizer:
library: yttm
tokenizer_model: null
vocab_file: null
special_tokens: null
vocab_size: null
encoder:
library: huggingface
model_name: bert-base-uncased
pretrained: false
decoder:
library: nemo
model_name: null
pretrained: false
max_sequence_length: 512
num_token_types: 2
embedding_dropout: 0.1
learn_positional_encodings: false
hidden_size: 512
inner_size: 2048
num_layers: 6
num_attention_heads: 8
ffn_dropout: 0.1
attn_score_dropout: 0.1
attn_layer_dropout: 0.1
hidden_act: relu
pre_ln: false
head:
num_layers: 1
activation: relu
log_softmax: true
dropout: 0.0
use_transformer_init: true
trainer:
gpus: 4
num_nodes: 1
max_epochs: 200
amp_level: O2 # O1/O2 for mixed precision
precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0
accelerator: ddp
checkpoint_callback: False
logger: False
log_every_n_steps: 50 # Interval of logging.
check_val_every_n_epoch: 1
exp_manager:
name: HuggingFaceEncoder
files_to_copy: []
tts/NeMo/examples/nlp/machine_translation/create_tarred_monolingual_dataset.py
+68
View File
@@ -0,0 +1,68 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='NMT dataset pre-processing')
parser.add_argument('--tokenizer_model', type=str, required=True, help='Path to tokenizer model')
parser.add_argument('--tokenizer_name', type=str, default='yttm', help='BPE Tokenizer Name, Options: [yttm]')
parser.add_argument('--bpe_droput', type=float, default=0.0, help='BPE dropout to use')
parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff')
parser.add_argument('--pkl_file_prefix', type=str, default='parallel', help='Prefix for tar and pickle files')
parser.add_argument('--fname', type=str, required=True, help='Path to monolingual data file')
parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models')
parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length')
parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length')
parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU')
parser.add_argument(
'--lines_per_dataset_fragment',
type=int,
default=1000000,
help='Number of lines to consider for bucketing and padding',
)
parser.add_argument(
'--num_batches_per_tarfile',
type=int,
default=1000,
help='Number of batches (pickle files) within each tarfile',
)
args = parser.parse_args()
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
if not os.path.exists(args.tokenizer_model):
assert FileNotFoundError("Could not find tokenizer model %s" % (args.tokenizer))
tokenizer_model = MTDataPreproc.get_monolingual_tokenizer(
tokenizer_name=args.tokenizer_name, tokenizer_model=args.tokenizer_model, bpe_dropout=args.bpe_droput
)
MTDataPreproc.preprocess_monolingual_dataset(
clean=args.clean,
fname=args.fname,
out_dir=args.out_dir,
tokenizer=tokenizer_model,
max_seq_length=args.max_seq_length,
min_seq_length=args.min_seq_length,
tokens_in_batch=args.tokens_in_batch,
lines_per_dataset_fragment=args.lines_per_dataset_fragment,
num_batches_per_tarfile=args.num_batches_per_tarfile,
pkl_file_prefix=args.pkl_file_prefix,
global_rank=0,
world_size=1,
)
tts/NeMo/examples/nlp/machine_translation/create_tarred_parallel_dataset.py
+143
View File
@@ -0,0 +1,143 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='NMT dataset pre-processing')
parser.add_argument('--shared_tokenizer', action="store_true", help='Whether to share encoder/decoder tokenizers')
parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff')
parser.add_argument('--tar_file_prefix', type=str, default='parallel', help='Prefix for tar files')
parser.add_argument('--bpe_dropout', type=float, default=0.1, help='Whether to share encoder/decoder tokenizers')
parser.add_argument('--src_fname', type=str, required=True, help='Path to the source file')
parser.add_argument('--tgt_fname', type=str, required=True, help='Path to the target file')
parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models')
parser.add_argument('--encoder_model_name', type=str, default=None, help='For use with pretrained encoders')
parser.add_argument(
'--decoder_model_name', type=str, default=None, help='For use with pretrained decoders (not yet supported)'
)
parser.add_argument(
'--encoder_tokenizer_model', type=str, default='None', help='Path to pre-trained encoder tokenizer model'
)
parser.add_argument(
'--encoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm]'
)
parser.add_argument('--encoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE')
parser.add_argument(
'--encoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE'
)
parser.add_argument('--encoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob')
parser.add_argument(
'--decoder_tokenizer_model', type=str, default='None', help='Path to pre-trained decoder tokenizer model'
)
parser.add_argument(
'--decoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm]'
)
parser.add_argument('--decoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE')
parser.add_argument(
'--decoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE'
)
parser.add_argument('--decoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob')
parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length')
parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length')
parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU')
parser.add_argument('--coverage', type=float, default=0.999, help='BPE character coverage [0-1]')
parser.add_argument(
'--lines_per_dataset_fragment',
type=int,
default=1000000,
help='Number of lines to consider for bucketing and padding',
)
parser.add_argument(
'--num_batches_per_tarfile',
type=int,
default=1000,
help='Number of batches (pickle files) within each tarfile',
)
parser.add_argument(
'--n_preproc_jobs', type=int, default=-2, help='Number of processes to use for creating the tarred dataset.',
)
args = parser.parse_args()
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
if (
args.encoder_tokenizer_model != 'None'
and args.decoder_tokenizer_model == 'None'
or args.decoder_tokenizer_model != 'None'
and args.encoder_tokenizer_model == 'None'
):
if args.shared_tokenizer:
raise ValueError(
'''
If using a pre-trained shared tokenizer,
both encoder and decoder tokenizers must be the same
'''
)
else:
raise ValueError('Both encoder and decoder pre-trained tokenizer models must be specified')
if args.encoder_tokenizer_model == 'None' and args.decoder_tokenizer_model == 'None':
encoder_tokenizer_model, decoder_tokenizer_model = MTDataPreproc.train_tokenizers(
out_dir=args.out_dir,
src_fname=args.src_fname,
tgt_fname=args.tgt_fname,
shared_tokenizer=args.shared_tokenizer,
encoder_tokenizer_name=args.encoder_tokenizer_name,
encoder_tokenizer_vocab_size=args.encoder_tokenizer_vocab_size,
encoder_tokenizer_coverage=args.encoder_tokenizer_coverage,
decoder_tokenizer_name=args.decoder_tokenizer_name,
decoder_tokenizer_vocab_size=args.decoder_tokenizer_vocab_size,
decoder_tokenizer_coverage=args.decoder_tokenizer_coverage,
global_rank=0,
)
else:
encoder_tokenizer_model, decoder_tokenizer_model = args.encoder_tokenizer_model, args.decoder_tokenizer_model
encoder_tokenizer, decoder_tokenizer = MTDataPreproc.get_enc_dec_tokenizers(
encoder_tokenizer_name=args.encoder_tokenizer_name,
encoder_tokenizer_model=encoder_tokenizer_model,
encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout,
decoder_tokenizer_name=args.decoder_tokenizer_name,
decoder_tokenizer_model=decoder_tokenizer_model,
decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout,
)
_, _ = MTDataPreproc.preprocess_parallel_dataset(
clean=args.clean,
src_fname=args.src_fname,
tgt_fname=args.tgt_fname,
out_dir=args.out_dir,
encoder_tokenizer_name=args.encoder_tokenizer_name,
encoder_model_name=args.encoder_model_name,
encoder_tokenizer_model=encoder_tokenizer_model,
encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout,
decoder_tokenizer_name=args.decoder_tokenizer_name,
decoder_model_name=args.decoder_model_name,
decoder_tokenizer_model=decoder_tokenizer_model,
decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout,
max_seq_length=args.max_seq_length,
min_seq_length=args.min_seq_length,
tokens_in_batch=args.tokens_in_batch,
lines_per_dataset_fragment=args.lines_per_dataset_fragment,
num_batches_per_tarfile=args.num_batches_per_tarfile,
tar_file_prefix=args.tar_file_prefix,
global_rank=0,
world_size=1,
n_jobs=args.n_preproc_jobs,
)
tts/NeMo/examples/nlp/machine_translation/enc_dec_nmt.py
+137
View File
@@ -0,0 +1,137 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
from typing import Optional
from omegaconf import OmegaConf
from pytorch_lightning import Trainer
from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc
from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import MTEncDecModelConfig
from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
from nemo.core.config import hydra_runner
from nemo.core.config.modelPT import NemoConfig
from nemo.core.config.pytorch_lightning import TrainerConfig
from nemo.utils import logging
from nemo.utils.config_utils import update_model_config
from nemo.utils.exp_manager import ExpManagerConfig, exp_manager
"""
Usage:
1. If you need to start docker and install NeMo, otherwise skip this step:
a. ```docker run --gpus all -it --rm -v /home/okuchaiev/repos/NeMo/:/NeMo -p 6006:6006 -v /mnt:/mnt --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:20.11-py3```
b. ```cd /NeMo```
c. ```./reinstall.sh```
2. Train a new tokenizer (or use pre-trained one):
```yttm bpe --data /mnt/D1/Data/NMT/wmt16_de_en/train.clean.en-de.shuffled.common --model tokenizer.BPE.8192.model --vocab_size 8192```
(To use WANDB, optionally, do login first)
``wandb login [YOUR WANDB login]``
3. Start training:
(This example for "base" model on 2 GPUs for 150000 steps with batch size of 12500 tokens per GPU)
python enc_dec_nmt.py \
--config-path=conf \
--config-name=aayn_base \
trainer.gpus=[0,1] \
~trainer.max_epochs \
+trainer.max_steps=150000 \
model.beam_size=4 \
model.max_generation_delta=5 \
model.label_smoothing=0.1 \
model.encoder_tokenizer.tokenizer_model=tokenizer.BPE.8192.model \
model.decoder_tokenizer.tokenizer_model=tokenizer.BPE.8192.model \
model.encoder.num_layers=6 \
model.encoder.hidden_size=512 \
model.encoder.inner_size=2048 \
model.encoder.num_attention_heads=8 \
model.encoder.ffn_dropout=0.1 \
model.decoder.num_layers=6 \
model.decoder.hidden_size=512 \
model.decoder.inner_size=2048 \
model.decoder.num_attention_heads=8 \
model.decoder.ffn_dropout=0.1 \
model.train_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/train.clean.de.shuffled \
model.train_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/train.clean.en.shuffled \
model.train_ds.tokens_in_batch=12500 \
model.validation_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.ref \
model.validation_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.src \
model.validation_ds.tokens_in_batch=8192 \
model.test_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.ref \
model.test_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.src \
model.optim.lr=0.001 \
model.optim.sched.warmup_ratio=0.05 \
+exp_manager.create_wandb_logger=True \
+exp_manager.wandb_logger_kwargs.name=TEST-nmt-base \
+exp_manager.wandb_logger_kwargs.project=nmt-de-en \
+exp_manager.create_checkpoint_callback=True \
+exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU \
+exp_manager.exp_dir=nmt_base \
+exp_manager.checkpoint_callback_params.mode=max
"""
@dataclass
class MTEncDecConfig(NemoConfig):
name: Optional[str] = 'MTEncDec'
do_training: bool = True
do_testing: bool = False
model: MTEncDecModelConfig = MTEncDecModelConfig()
trainer: Optional[TrainerConfig] = TrainerConfig()
exp_manager: Optional[ExpManagerConfig] = ExpManagerConfig(name='MTEncDec', files_to_copy=[])
@hydra_runner(config_path="conf", config_name="aayn_base")
def main(cfg: MTEncDecConfig) -> None:
# merge default config with user specified config
default_cfg = MTEncDecConfig()
cfg = update_model_config(default_cfg, cfg)
logging.info("\n\n************** Experiment configuration ***********")
logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
# training is managed by PyTorch Lightning
trainer = Trainer(**cfg.trainer)
# tokenizers will be trained and and tarred training data will be created if needed
# model config is then updated
if cfg.model.preproc_out_dir is not None:
MTDataPreproc(cfg=cfg.model, trainer=trainer)
# experiment logs, checkpoints, and auto-resume are managed by exp_manager and PyTorch Lightning
exp_manager(trainer, cfg.exp_manager)
# everything needed to train translation models is encapsulated in the NeMo MTEncdDecModel
mt_model = MTEncDecModel(cfg.model, trainer=trainer)
logging.info("\n\n************** Model parameters and their sizes ***********")
for name, param in mt_model.named_parameters():
print(name, param.size())
logging.info("***********************************************************\n\n")
if cfg.do_training:
trainer.fit(mt_model)
if cfg.do_testing:
trainer.test(mt_model)
if __name__ == '__main__':
main()
tts/NeMo/examples/nlp/machine_translation/nmt_transformer_infer.py
+122
View File
@@ -0,0 +1,122 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Given NMT model's .nemo file, this script can be used to translate text.
USAGE Example:
1. Obtain text file in src language. You can use sacrebleu to obtain standard test sets like so:
sacrebleu -t wmt14 -l de-en --echo src > wmt14-de-en.src
2. Translate:
python nmt_transformer_infer.py --model=[Path to .nemo file] --srctext=wmt14-de-en.src --tgtout=wmt14-de-en.pre
"""
from argparse import ArgumentParser
import torch
import nemo.collections.nlp as nemo_nlp
from nemo.collections.nlp.modules.common.transformer import (
BeamSearchSequenceGenerator,
BeamSearchSequenceGeneratorWithLanguageModel,
)
from nemo.utils import logging
def main():
parser = ArgumentParser()
parser.add_argument("--model", type=str, required=True, help="")
parser.add_argument("--srctext", type=str, required=True, help="")
parser.add_argument("--tgtout", type=str, required=True, help="")
parser.add_argument("--batch_size", type=int, default=256, help="")
parser.add_argument("--beam_size", type=int, default=4, help="")
parser.add_argument("--len_pen", type=float, default=0.6, help="")
parser.add_argument("--max_delta_length", type=int, default=5, help="")
parser.add_argument("--target_lang", type=str, default=None, help="")
parser.add_argument("--source_lang", type=str, default=None, help="")
# shallow fusion specific parameters
parser.add_argument("--lm_model", type=str, default=None, help="")
parser.add_argument("--fusion_coef", type=float, default=0.0, help="")
args = parser.parse_args()
torch.set_grad_enabled(False)
if args.model.endswith(".nemo"):
logging.info("Attempting to initialize from .nemo file")
model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(restore_path=args.model)
src_text = []
tgt_text = []
else:
raise NotImplemented(f"Only support .nemo files, but got: {args.model}")
if torch.cuda.is_available():
model = model.cuda()
if args.lm_model is not None:
lm_model = nemo_nlp.models.language_modeling.TransformerLMModel.restore_from(restore_path=args.lm_model).eval()
model.beam_search = BeamSearchSequenceGeneratorWithLanguageModel(
embedding=model.decoder.embedding,
decoder=model.decoder.decoder,
log_softmax=model.log_softmax,
bos=model.decoder_tokenizer.bos_id,
pad=model.decoder_tokenizer.pad_id,
eos=model.decoder_tokenizer.eos_id,
language_model=lm_model,
fusion_coef=args.fusion_coef,
max_sequence_length=model.decoder.max_sequence_length,
beam_size=args.beam_size,
len_pen=args.len_pen,
max_delta_length=args.max_delta_length,
)
else:
model.beam_search = BeamSearchSequenceGenerator(
embedding=model.decoder.embedding,
decoder=model.decoder.decoder,
log_softmax=model.log_softmax,
bos=model.decoder_tokenizer.bos_id,
pad=model.decoder_tokenizer.pad_id,
eos=model.decoder_tokenizer.eos_id,
max_sequence_length=model.decoder.max_sequence_length,
beam_size=args.beam_size,
len_pen=args.len_pen,
max_delta_length=args.max_delta_length,
)
logging.info(f"Translating: {args.srctext}")
count = 0
with open(args.srctext, 'r') as src_f:
for line in src_f:
src_text.append(line.strip())
if len(src_text) == args.batch_size:
res = model.translate(text=src_text, source_lang=args.source_lang, target_lang=args.target_lang)
if len(res) != len(src_text):
print(len(res))
print(len(src_text))
print(res)
print(src_text)
tgt_text += res
src_text = []
count += 1
# if count % 300 == 0:
# print(f"Translated {count} sentences")
if len(src_text) > 0:
tgt_text += model.translate(text=src_text, source_lang=args.source_lang, target_lang=args.target_lang)
with open(args.tgtout, 'w') as tgt_f:
for line in tgt_text:
tgt_f.write(line + "\n")
if __name__ == '__main__':
main() # noqa pylint: disable=no-value-for-parameter
tts/NeMo/examples/nlp/machine_translation/nmt_webapp/README.rst
+7
View File
@@ -0,0 +1,7 @@
**STEPS 2 RUN**
===============
0. (Make sure you have Flask installed - ``pip install -r requirements.txt``)
1. Edit "config.json" file to only contain models you need. If model's location starts with "NGC/" - it will load this model from NVIDIA's NGC. Otherwise, specify full path to .nemo file.
4. To run: ``python nmt_service.py``
5. To translate: ``http://127.0.0.1:5000/translate?text=Frohe%20Weihnachten`` (here %20 means space)
6. Run web UI: ``python -m http.server``
tts/NeMo/examples/nlp/machine_translation/nmt_webapp/config.json
+12
View File
@@ -0,0 +1,12 @@
{
"en-de": "NGC/nmt_en_de_transformer12x2",
"de-en": "NGC/nmt_de_en_transformer12x2",
"en-es": "NGC/nmt_en_es_transformer12x2",
"es-en": "NGC/nmt_es_en_transformer12x2",
"en-ru": "NGC/nmt_en_ru_transformer6x6",
"ru-en": "NGC/nmt_ru_en_transformer6x6",
"en-fr": "NGC/nmt_en_fr_transformer12x2",
"fr-en": "NGC/nmt_fr_en_transformer12x2",
"en-zh": "NGC/nmt_en_zh_transformer6x6",
"zh-en": "NGC/nmt_zh_en_transformer6x6"
}
tts/NeMo/examples/nlp/machine_translation/nmt_webapp/index.html
+140
View File
@@ -0,0 +1,140 @@
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Neural Machine Translation Demo App</title>
<link rel="stylesheet" href="//code.jquery.com/ui/1.12.1/themes/base/jquery-ui.css">
<link rel="stylesheet" href="/resources/demos/style.css">
<style>
.ui-resizable-se {
bottom: 17px;
}
</style>
<link rel="stylesheet" href="style.css">
<script src="https://code.jquery.com/jquery-1.12.4.js"></script>
<script src="https://code.jquery.com/ui/1.12.1/jquery-ui.js"></script>
<script>
$( function() {
$( "#tgt_text" ).resizable({
handles: "se"
});
} );
$( function() {
$( "#src_text" ).resizable({
handles: "se"
});
} );
$( function() {
$( "#translate_button" ).button();
$( "#translate_button" ).click( function( event ) {
$.ajax({
url: "http://10.110.40.187:5000/translate",
type: "GET",
data: {
text: $( "#src_text" ).val(),
langpair: $('input[type=radio][name=langpair]:checked').attr('id'),
do_moses: true
},
success: function (result) {
$( "#tgt_text" ).val(result.translation);
},
error: function (err) {
alert('Sorry, can not reach translation service');
},
})
} );
} );
</script>
</head>
<body>
<div class="contact-us">
<h1>Neural Machine Translation Demo</h1>
<p>
<br/>
<h4>Running NeMo commit id: v1.0.2 </h4>
<br/>
</p>
<p>
<fieldset>
<legend>Select translation direction: </legend>
<table>
<tr>
<td>
<label for="en-de">English to German</label>
<input type="radio" name="langpair" id="en-de">
</td>
<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
<td>
<label for="de-en">German to English</label>
<input type="radio" name="langpair" id="de-en">
</td>
<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
<td>
<label for="es-en">Spanish to English</label>
<input type="radio" name="langpair" id="es-en">
</td>
<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
<td>
<label for="en-es" selected="selected">English to Spanish</label>
<input type="radio" name="langpair" id="en-es" checked>
</td>
</tr>
<tr>
<td>
<label for="en-ru">English to Russian</label>
<input type="radio" name="langpair" id="en-ru">
</td>
<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
<td>
<label for="ru-en">Russian to English</label>
<input type="radio" name="langpair" id="ru-en">
</td>
<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
<td>
<label for="en-fr">English to French</label>
<input type="radio" name="langpair" id="en-fr">
</td>
<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
<td>
<label for="fr-en">French to English</label>
<input type="radio" name="langpair" id="fr-en">
</td>
<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
<td>
<label for="en-zh">English to Chinese</label>
<input type="radio" name="langpair" id="en-zh">
</td>
<td>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</td>
<td>
<label for="zh-en">Chinese to English</label>
<input type="radio" name="langpair" id="zh-en">
</td>
</tr>
</table>
</fieldset>
<div>
<table>
<tr>
<td>
<textarea id="src_text" rows="12" cols="60">Type something here, chose language pair and click "Translate" button.</textarea>
</td>
</tr>
<tr>
<td>
<button id="translate_button">Translate !</button>
</td>
</tr>
<tr>
<td>
<textarea id="tgt_text" rows="12" cols="60" readonly></textarea>
</td>
</tr>
</table>
</div>
</p>
</div>
</body>
</html>
tts/NeMo/examples/nlp/machine_translation/nmt_webapp/nmt_service.py
+100
View File
@@ -0,0 +1,100 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import time
import flask
import torch
from flask import Flask, json, request
from flask_cors import CORS
import nemo.collections.nlp as nemo_nlp
from nemo.utils import logging
MODELS_DICT = {}
model = None
api = Flask(__name__)
CORS(api)
def initialize(config_file_path: str):
"""
Loads 'language-pair to NMT model mapping'
"""
__MODELS_DICT = None
logging.info("Starting NMT service")
logging.info(f"I will attempt to load all the models listed in {config_file_path}.")
logging.info(f"Edit {config_file_path} to disable models you don't need.")
if torch.cuda.is_available():
logging.info("CUDA is available. Running on GPU")
else:
logging.info("CUDA is not available. Defaulting to CPUs")
# read config
with open(config_file_path) as f:
__MODELS_DICT = json.load(f)
if __MODELS_DICT is not None:
for key, value in __MODELS_DICT.items():
logging.info(f"Loading model for {key} from file: {value}")
if value.startswith("NGC/"):
model = nemo_nlp.models.machine_translation.MTEncDecModel.from_pretrained(model_name=value[4:])
else:
model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(restore_path=value)
if torch.cuda.is_available():
model = model.cuda()
MODELS_DICT[key] = model
else:
raise ValueError("Did not find the config.json or it was empty")
logging.info("NMT service started")
@api.route('/translate', methods=['GET', 'POST', 'OPTIONS'])
def get_translation():
try:
time_s = time.time()
langpair = request.args["langpair"]
src = request.args["text"]
do_moses = request.args.get('do_moses', False)
if langpair in MODELS_DICT:
if do_moses:
result = MODELS_DICT[langpair].translate(
[src], source_lang=langpair.split('-')[0], target_lang=langpair.split('-')[1]
)
else:
result = MODELS_DICT[langpair].translate([src])
duration = time.time() - time_s
logging.info(
f"Translated in {duration}. Input was: {request.args['text']} <############> Translation was: {result[0]}"
)
res = {'translation': result[0]}
response = flask.jsonify(res)
response.headers.add('Access-Control-Allow-Origin', '*')
return response
else:
logging.error(f"Got the following langpair: {langpair} which was not found")
except Exception as ex:
res = {'translation': str(ex)}
response = flask.jsonify(res)
response.headers.add('Access-Control-Allow-Origin', '*')
return res
if __name__ == '__main__':
initialize('config.json')
api.run(host='0.0.0.0')
tts/NeMo/examples/nlp/machine_translation/nmt_webapp/requirements.txt
+3
View File
@@ -0,0 +1,3 @@
nemo_toolkit[all]>=1.0.0rc1
flask
flask_cors
tts/NeMo/examples/nlp/machine_translation/nmt_webapp/style.css
+111
View File
@@ -0,0 +1,111 @@
@import url("https://fonts.googleapis.com/css?family=Spartan&display=swap");
* {
margin: 0;
padding: 0;
}
body {
/*background: url("https://img.wallpapersafari.com/desktop/1680/1050/65/18/8AYkSt.jpg") center center;*/
background: url("https://images5.alphacoders.com/550/550948.jpg") center center;
background-size: cover;
width: 100vw;
height: 100vh;
display: grid;
align-items: center;
justify-items: center;
font-size: 12pt;
font-family: "Spartan";
color: #ffffff;
}
.contact-us {
/*background: #1f1d1d;*/
padding: 50px 100px;
border-top: 10px solid #229955;
}
label, input, textarea {
display: block;
width: 100%;
font-size: 12pt;
line-height: 24pt;
font-family: "Spartan";
}
input {
margin-bottom: 24pt;
}
h1 {
color: #ffffff;
font-weight: bold;
font-size: 22pt;
line-height: 24pt;
font-style: italic;
margin: 0 0 0.5em 0;
}
span {
font-size: 8pt;
}
em {
color: #229955;
font-weight: bold;
}
input, textarea {
border: none;
border: 1px solid rgba(0, 0, 0, 0.1);
border-radius: 2px;
background: #f8f4e5;
padding-left: 5px;
outline: none;
}
input:focus, textarea:focus {
border: 1px solid #6bd4b1;
}
textarea {
resize: none;
}
fieldset {
overflow: hidden
}
button {
display: block;
float: left;
line-height: 24pt;
padding: 0 20px;
border: none;
background: #f45702;
color: white;
letter-spacing: 2px;
transition: 0.2s all ease-in-out;
border-bottom: 2px solid transparent;
outline: none;
}
button:hover {
background: inherit;
color: #ea6153;
border-bottom: 2px solid #f45702;
}
::selection {
background: #ffc7b8;
}
input:-webkit-autofill,
input:-webkit-autofill:hover,
input:-webkit-autofill:focus,
textarea:-webkit-autofill,
textarea:-webkit-autofill:hover,
textarea:-webkit-autofill:focus {
border: 1px solid #6bd4b1;
-webkit-text-fill-color: #2A293E;
-webkit-box-shadow: 0 0 0px 1000px #f8f4e5 inset;
transition: background-color 5000s ease-in-out 0s;
}
tts/NeMo/examples/nlp/machine_translation/translate_ddp.py
+122
View File
@@ -0,0 +1,122 @@
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from argparse import ArgumentParser
import torch
import torch.multiprocessing as mp
from torch.utils.data import DataLoader
from nemo.collections.nlp.data.language_modeling import TarredSentenceDataset
from nemo.collections.nlp.data.machine_translation import TarredTranslationDataset
from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
from nemo.utils import logging
def get_args():
parser = ArgumentParser(description='Batch translation of sentences from a pre-trained model on multiple GPUs')
parser.add_argument("--model", type=str, required=True, help="Path to the .nemo translation model file")
parser.add_argument(
"--text2translate", type=str, required=True, help="Path to the pre-processed tarfiles for translation"
)
parser.add_argument("--result_dir", type=str, required=True, help="Folder to write translation results")
parser.add_argument(
"--twoside", action="store_true", help="Set flag when translating the source side of a parallel dataset"
)
parser.add_argument(
'--metadata_path', type=str, required=True, help="Path to the JSON file that contains dataset info"
)
parser.add_argument('--topk', type=int, default=500, help="Value of k for topk sampling")
parser.add_argument('--src_language', type=str, required=True, help="Source lang ID for detokenization")
parser.add_argument('--tgt_language', type=str, required=True, help="Target lang ID for detokenization")
parser.add_argument(
'--reverse_lang_direction',
action="store_true",
help="Reverse source and target language direction for parallel dataset",
)
parser.add_argument('--n_gpus', type=int, default=-1, help="Number of GPUs to use")
args = parser.parse_args()
return args
def translate(rank, world_size, args):
if args.model.endswith(".nemo"):
logging.info("Attempting to initialize from .nemo file")
model = MTEncDecModel.restore_from(restore_path=args.model, map_location=f"cuda:{rank}")
elif args.model.endswith(".ckpt"):
logging.info("Attempting to initialize from .ckpt file")
model = MTEncDecModel.load_from_checkpoint(checkpoint_path=args.model, map_location=f"cuda:{rank}")
model.replace_beam_with_sampling(topk=args.topk)
model.eval()
if args.twoside:
dataset = TarredTranslationDataset(
text_tar_filepaths=args.text2translate,
metadata_path=args.metadata_path,
encoder_tokenizer=model.encoder_tokenizer,
decoder_tokenizer=model.decoder_tokenizer,
shuffle_n=100,
shard_strategy="scatter",
world_size=world_size,
global_rank=rank,
reverse_lang_direction=args.reverse_lang_direction,
)
else:
dataset = TarredSentenceDataset(
text_tar_filepaths=args.text2translate,
metadata_path=args.metadata_path,
tokenizer=model.encoder_tokenizer,
shuffle_n=100,
shard_strategy="scatter",
world_size=world_size,
global_rank=rank,
)
loader = DataLoader(dataset, batch_size=1)
result_dir = os.path.join(args.result_dir, f'rank{rank}')
os.makedirs(result_dir, exist_ok=True)
originals_file_name = os.path.join(result_dir, 'originals.txt')
translations_file_name = os.path.join(result_dir, 'translations.txt')
num_translated_sentences = 0
with open(originals_file_name, 'w') as of, open(translations_file_name, 'w') as tf:
for batch_idx, batch in enumerate(loader):
for i in range(len(batch)):
if batch[i].ndim == 3:
batch[i] = batch[i].squeeze(dim=0)
batch[i] = batch[i].to(rank)
if args.twoside:
src_ids, src_mask, _, _, _ = batch
else:
src_ids, src_mask = batch
if batch_idx % 100 == 0:
logging.info(
f"{batch_idx} batches ({num_translated_sentences} sentences) were translated by process with "
f"rank {rank}"
)
num_translated_sentences += len(src_ids)
inputs, translations = model.batch_translate(src=src_ids, src_mask=src_mask)
for src, translation in zip(inputs, translations):
of.write(src + '\n')
tf.write(translation + '\n')
def main() -> None:
args = get_args()
world_size = torch.cuda.device_count() if args.n_gpus == -1 else args.n_gpus
mp.spawn(translate, args=(world_size, args), nprocs=world_size, join=True)
if __name__ == '__main__':
main()