mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Brandon Thomas
113 lines
5.2 KiB
Python
113 lines
5.2 KiB
Python
# Copyright (c) 2020, MeetKai Inc. All rights reserved.
|
|
# 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.
|
|
|
|
"""
|
|
This script contains an example on how to train and save a NeuralMachineTranslationModel.
|
|
NeuralMachineTranslationModel in NeMo supports sequence to sequence problems such as language translation
|
|
and text summarization, provided the data follows the format specified below.
|
|
|
|
|
|
***Data format***
|
|
NeuralMachineTranslationModel requires the data to be stored in TAB separated files (.tsv) with two columns of
|
|
sentence and label, where the first line is a header of format:
|
|
sentence[TAB]label
|
|
And each line is of the format:
|
|
[SENTENCE][TAB][LABEL]
|
|
|
|
If your dataset is stored in another format, you need to convert it to this format to use a
|
|
NeuralMachineTranslationModel.
|
|
|
|
|
|
***Setting the configs***
|
|
This script uses the `/examples/nlp/text2sparql/conf/text2sparql_config.yaml` config file by default.
|
|
You may update the config file from the file directly or by using the command line arguments.
|
|
Another other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'.
|
|
|
|
A NeuralMachineTranslationModel's config file declares multiple import sections. They are:
|
|
- trainer: Arguments to be passed to PyTorch Lightning.
|
|
- model: All arguments that relate to the Model - language_model, tokenizers, datasets, optimizer, generate.
|
|
- exp_manager: Arguments to be passed to NeMo's experiment manager.
|
|
- hydra: Arguments to be passed to Hydra.
|
|
|
|
If using text2sparql_config.yaml, you must first update the following fields in the config:
|
|
- model.nemo_path: Model save path. Eg. [PATH]/bart.nemo
|
|
- model.data_dir: Path to data directory. Alternatively, you can adjust the file paths directly:
|
|
- model.train_ds.filepath
|
|
- model.validation_ds.filepath
|
|
- model.test_ds.filepath
|
|
- exp_manager.exp_dir: Directory to log results from the experiment.
|
|
|
|
It is highly recommended to also adjust these parameters as necessary:
|
|
- trainer.gpus: Set to 0 to use CPU. Otherwise the number denotes the number of GPUs.
|
|
- trainer.max_epochs: Maximum number of epochs to train for.
|
|
- model.batch_size: 8 is sufficient to train a decent Bart model for Text2Sparql.
|
|
- model.max_seq_length: Maximum (tokenized) sequence length. 150 is sufficient for Text2Sparql.
|
|
- model.language_model.pretrained_model_name: End2end pretrained model name from huggingface.
|
|
- model.encoder_tokenizer.tokenizer_name: Pretrained tokenizer name from huggingface.
|
|
- model.decoder_tokenizer.tokenizer_name: The same as above, as the tokenizer will handle encoding and decoding.
|
|
- model.optim.lr: Learning rate.
|
|
|
|
You can also specify an encoder and decoder rather than using an end2end model like Bart by defining these parameters:
|
|
- model.language_model.pretrained_encoder_model_name: Pretrained huggingface encoder model name.
|
|
- model.encoder_tokenizer.tokenizer_name: Pretrained huggingface encoder tokenizer name.
|
|
- model.language_model.pretrained_decoder_model_name: Pretrained huggingface decoder model name.
|
|
- model.decoder_tokenizer.tokenizer_name: Pretrained huggingface decoder tokenizer name.
|
|
- model.language_model.pretrained_model_name: Set this to null.
|
|
|
|
|
|
***How to run the script?***
|
|
- First, download the data to TGT_DATA_DIR (see ./data/import_datasets.py):
|
|
|
|
SRC_DATA_DIR=./data/text2sparql_src
|
|
TGT_DATA_DIR=./data/text2sparql_tgt
|
|
|
|
python ./data/import_datasets.py \
|
|
--source_data_dir $SRC_DATA_DIR \
|
|
--target_data_dir $TGT_DATA_DIR
|
|
|
|
- And run the following to train and save the model:
|
|
|
|
python text2sparql.py \
|
|
model.train_ds.filepath="$TGT_DATA_DIR"/train.tsv \
|
|
model.validation_ds.filepath="$TGT_DATA_DIR"/test_easy.tsv \
|
|
model.test_ds.filepath="$TGT_DATA_DIR"/test_hard.tsv \
|
|
model.batch_size=16 \
|
|
model.nemo_path=./NeMo_logs/bart.nemo \
|
|
exp_manager.exp_dir=./NeMo_logs
|
|
"""
|
|
|
|
import pytorch_lightning as pl
|
|
from omegaconf import DictConfig, OmegaConf
|
|
|
|
from nemo.collections.nlp.models.neural_machine_translation import NeuralMachineTranslationModel
|
|
from nemo.core.config import hydra_runner
|
|
from nemo.utils import logging
|
|
from nemo.utils.exp_manager import exp_manager
|
|
|
|
|
|
@hydra_runner(config_path="conf", config_name="text2sparql_config")
|
|
def main(cfg: DictConfig) -> None:
|
|
logging.info(f"Config:\n {OmegaConf.to_yaml(cfg)}")
|
|
trainer = pl.Trainer(**cfg.trainer)
|
|
exp_manager(trainer, cfg.get("exp_manager", None))
|
|
nmt_model = NeuralMachineTranslationModel(cfg.model, trainer=trainer)
|
|
trainer.fit(nmt_model)
|
|
if cfg.model.nemo_path:
|
|
nmt_model.save_to(cfg.model.nemo_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|