diff --git a/docs/source/nlp/models.rst b/docs/source/nlp/models.rst
index c6b9adbad02a..7218525b29e7 100755
--- a/docs/source/nlp/models.rst
+++ b/docs/source/nlp/models.rst
@@ -20,3 +20,4 @@ NeMo's NLP collection supports the following models:
information_retrieval
nlp_model
machine_translation
+ text_normalization
diff --git a/docs/source/nlp/nlp_all.bib b/docs/source/nlp/nlp_all.bib
index f291310941e4..a7084fe78159 100644
--- a/docs/source/nlp/nlp_all.bib
+++ b/docs/source/nlp/nlp_all.bib
@@ -100,10 +100,26 @@ @article{post2018call
}
@misc{zhang2021sgdqa,
- title={SGD-QA: Fast Schema-Guided Dialogue State Tracking for Unseen Services},
+ title={SGD-QA: Fast Schema-Guided Dialogue State Tracking for Unseen Services},
author={Yang Zhang and Vahid Noroozi and Evelina Bakhturina and Boris Ginsburg},
year={2021},
eprint={2105.08049},
archivePrefix={arXiv},
primaryClass={cs.CL}
-}
\ No newline at end of file
+}
+
+@article{Sproat2016RNNAT,
+ title={RNN Approaches to Text Normalization: A Challenge},
+ author={R. Sproat and Navdeep Jaitly},
+ journal={ArXiv},
+ year={2016},
+ volume={abs/1611.00068}
+}
+
+@article{Zhang2019NeuralMO,
+ title={Neural Models of Text Normalization for Speech Applications},
+ author={Hao Zhang and R. Sproat and Axel H. Ng and Felix Stahlberg and Xiaochang Peng and Kyle Gorman and B. Roark},
+ journal={Computational Linguistics},
+ year={2019},
+ pages={293-338}
+}
diff --git a/docs/source/nlp/text_normalization.rst b/docs/source/nlp/text_normalization.rst
new file mode 100644
index 000000000000..f3871c3115ca
--- /dev/null
+++ b/docs/source/nlp/text_normalization.rst
@@ -0,0 +1,159 @@
+.. _text_normalization:
+
+Text Normalization Models
+==========================
+Text normalization is the task of converting a written text into its spoken form. For example,
+``$123`` should be verbalized as ``one hundred twenty three dollars``, while ``123 King Ave``
+should be verbalized as ``one twenty three King Avenue``. At the same time, the inverse problem
+is about converting a spoken sequence (e.g., an ASR output) into its written form.
+
+NeMo has an implementation that allows you to build a neural-based system that is able to do
+both text normalization (TN) and also inverse text normalization (ITN). At a high level, the
+system consists of two individual components:
+
+- `DuplexTaggerModel `__ - a Transformer-based tagger for identifying "semiotic" spans in the input (e.g., spans that are about times, dates, or monetary amounts).
+- `DuplexDecoderModel `__ - a Transformer-based seq2seq model for decoding the semiotic spans into their appropriate forms (e.g., spoken forms for TN and written forms for ITN).
+
+The typical workflow is to first train a DuplexTaggerModel and also a DuplexDecoderModel. An example training script
+is provided: `duplex_text_normalization_train.py `__.
+After that, the two trained models can be used to initialize a `DuplexTextNormalizationModel `__ that can be used for end-to-end inference.
+An example script for evaluation and inference is provided: `duplex_text_normalization_test.py `__. The term
+*duplex* refers to the fact that our system can be trained to do both TN and ITN. However, you can also specifically train the system for only one of the tasks.
+
+NeMo Data Format
+-----------
+Both the DuplexTaggerModel model and the DuplexDecoderModel model use the same simple text format
+as the dataset. The data needs to be stored in TAB separated files (``.tsv``) with three columns.
+The first of which is the "semiotic class" (e.g., numbers, times, dates) , the second is the token
+in written form, and the third is the spoken form. An example sentence in the dataset is shown below.
+In the example, ``sil`` denotes that a token is a punctuation while ``self`` denotes that the spoken form is the
+same as the written form. It is expected that a complete dataset contains three files: ``train.tsv``, ``dev.tsv``,
+and ``test.tsv``.
+
+.. code::
+
+ PLAIN The
+ PLAIN company 's
+ PLAIN revenues
+ PLAIN grew
+ PLAIN four
+ PLAIN fold
+ PLAIN between
+ DATE 2005 two thousand five
+ PLAIN and
+ DATE 2008 two thousand eight
+ PUNCT . sil
+
+
+
+An example script for generating a dataset in this format from the `Google text normalization dataset `_
+can be found at `NeMo/examples/nlp/duplex_text_normalization/google_data_preprocessing.py `__.
+Note that the script also does some preprocessing on the spoken forms of the URLs. For example,
+given the URL "Zimbio.com", the original expected spoken form in the Google dataset is
+"z_letter i_letter m_letter b_letter i_letter o_letter dot c_letter o_letter m_letter".
+However, our script will return a more concise output which is "zim bio dot com".
+
+More information about the Google text normalization dataset can be found in the paper `RNN Approaches to Text Normalization: A Challenge `__ :cite:`nlp-textnorm-Sproat2016RNNAT`.
+
+
+Model Training
+--------------
+
+An example training script is provided: `duplex_text_normalization_train.py `__.
+The config file used for the example is at `duplex_tn_config.yaml `__.
+You can change any of these parameters directly from the config file or update them with the command-line arguments.
+
+The config file contains three main sections. The first section contains the configs for the tagger, the second section is about the decoder,
+and the last section is about the dataset. Most arguments in the example config file are quite self-explanatory (e.g.,
+*decoder_model.optim.lr* refers to the learning rate for training the decoder). We have set most of the hyper-parameters to
+be the values that we found to be effective. Some arguments that you may want to modify are:
+
+- *data.base_dir*: The path to the dataset directory. It is expected that the directory contains three files: train.tsv, dev.tsv, and test.tsv.
+
+- *tagger_model.nemo_path*: This is the path where the final trained tagger model will be saved to.
+
+- *decoder_model.nemo_path*: This is the path where the final trained decoder model will be saved to.
+
+Example of a training command:
+
+.. code::
+
+ python examples/nlp/duplex_text_normalization/duplex_text_normalization_train.py \
+ data.base_dir= \
+ mode={tn,itn,joint}
+
+There are 3 different modes. "tn" mode is for training a system for TN only.
+"itn" mode is for training a system for ITN. "joint" is for training a system
+that can do both TN and ITN at the same time. Note that the above command will
+first train a tagger and then train a decoder sequentially.
+
+You can also train only a tagger (without training a decoder) by running the
+following command:
+
+.. code::
+
+ python examples/nlp/duplex_text_normalization/duplex_text_normalization_train.py \
+ data.base_dir=PATH_TO_DATASET_DIR \
+ mode={tn,itn,joint} \
+ decoder_model.do_training=false
+
+Or you can also train only a decoder (without training a tagger):
+
+.. code::
+
+ python examples/nlp/duplex_text_normalization/duplex_text_normalization_train.py \
+ data.base_dir=PATH_TO_DATASET_DIR \
+ mode={tn,itn,joint} \
+ tagger_model.do_training=false
+
+
+Model Architecture
+--------------
+
+The tagger model first uses a Transformer encoder (e.g., DistilRoBERTa) to build a
+contextualized representation for each input token. It then uses a classification head
+to predict the tag for each token (e.g., if a token should stay the same, its tag should
+be ``SAME``). The decoder model then takes the semiotic spans identified by the tagger and
+transform them into the appropriate forms (e.g., spoken forms for TN and written forms for ITN).
+The decoder model is essentially a Transformer-based encoder-decoder seq2seq model (e.g., the example
+training script uses the T5-base model by default). Overall, our design is partly inspired by the
+RNN-based sliding window model proposed in the paper
+`Neural Models of Text Normalization for Speech Applications `__ :cite:`nlp-textnorm-Zhang2019NeuralMO`.
+
+We introduce a simple but effective technique to allow our model to be duplex. Depending on the
+task the model is handling, we append the appropriate prefix to the input. For example, suppose
+we want to transform the text ``I live in 123 King Ave`` to its spoken form (i.e., TN problem),
+then we will simply append the prefix ``tn`` to it and so the final input to our models will actually
+be ``tn I live in tn 123 King Ave``. Similarly, for the ITN problem, we just append the prefix ``itn``
+to the input.
+
+To improve the effectiveness and robustness of our models, we also apply some simple data
+augmentation techniques during training.
+
+Data Augmentation for Training DuplexTaggerModel
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+In the Google English TN training data, about 93% of the tokens are not in any semiotic span. In other words, the ground-truth tags of most tokens are of trivial types (i.e., ``SAME`` and ``PUNCT``). To alleviate this class imbalance problem,
+for each original instance with several semiotic spans, we create a new instance by simply concatenating all the semiotic spans together. For example, considering the following ITN instance:
+
+Original instance: ``[The|SAME] [revenues|SAME] [grew|SAME] [a|SAME] [lot|SAME] [between|SAME] [two|B-TRANSFORM] [thousand|I-TRANSFORM] [two|I-TRANSFORM] [and|SAME] [two|B-TRANSFORM] [thousand|I-TRANSFORM] [five|I-TRANSFORM] [.|PUNCT]``
+
+Augmented instance: ``[two|B-TRANSFORM] [thousand|I-TRANSFORM] [two|I-TRANSFORM] [two|B-TRANSFORM] [thousand|I-TRANSFORM] [five|I-TRANSFORM]``
+
+The argument ``data.train_ds.tagger_data_augmentation`` in the config file controls whether this data augmentation will be enabled or not.
+
+
+Data Augmentation for Training DuplexDecoderModel
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Since the tagger may not be perfect, the inputs to the decoder may not all be semiotic spans. Therefore, to make the decoder become more robust against the tagger's potential errors,
+we train the decoder with not only semiotic spans but also with some other more "noisy" spans. This way even if the tagger makes some errors, there will still be some chance that the
+final output is still correct.
+
+The argument ``data.train_ds.decoder_data_augmentation`` in the config file controls whether this data augmentation will be enabled or not.
+
+References
+----------
+
+.. bibliography:: nlp_all.bib
+ :style: plain
+ :labelprefix: NLP-TEXTNORM
+ :keyprefix: nlp-textnorm-
diff --git a/examples/nlp/duplex_text_normalization/conf/duplex_tn_config.yaml b/examples/nlp/duplex_text_normalization/conf/duplex_tn_config.yaml
new file mode 100644
index 000000000000..c9aa55f850c3
--- /dev/null
+++ b/examples/nlp/duplex_text_normalization/conf/duplex_tn_config.yaml
@@ -0,0 +1,136 @@
+name: &name DuplexTextNormalization
+mode: joint # Three possible choices ['tn', 'itn', 'joint']
+
+# Pretrained Nemo Models
+tagger_pretrained_model: null
+decoder_pretrained_model: null
+
+# Tagger
+tagger_trainer:
+ gpus: 1 # the number of gpus, 0 for CPU
+ num_nodes: 1
+ max_epochs: 5 # the number of training epochs
+ checkpoint_callback: false # provided by exp_manager
+ logger: false # provided by exp_manager
+ accumulate_grad_batches: 1 # accumulates grads every k batches
+ gradient_clip_val: 0.0
+ amp_level: O0 # O1/O2 for mixed precision
+ precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+ accelerator: ddp
+
+tagger_model:
+ do_training: true
+ transformer: distilroberta-base
+ tokenizer: ${tagger_model.transformer}
+ nemo_path: ${tagger_exp_manager.exp_dir}/tagger_model.nemo # exported .nemo path
+
+ optim:
+ name: adamw
+ lr: 5e-5
+ weight_decay: 0.01
+
+ sched:
+ name: WarmupAnnealing
+
+ # pytorch lightning args
+ monitor: val_token_precision
+ reduce_on_plateau: false
+
+ # scheduler config override
+ warmup_steps: null
+ warmup_ratio: 0.1
+ last_epoch: -1
+
+tagger_exp_manager:
+ exp_dir: exps # where to store logs and checkpoints
+ name: tagger_training # name of experiment
+ create_tensorboard_logger: True
+ create_checkpoint_callback: True
+ checkpoint_callback_params:
+ save_top_k: 3
+ monitor: "val_token_precision"
+ mode: "max"
+ save_best_model: true
+ always_save_nemo: true
+
+# Decoder
+decoder_trainer:
+ gpus: 1 # the number of gpus, 0 for CPU
+ num_nodes: 1
+ max_epochs: 3 # the number of training epochs
+ checkpoint_callback: false # provided by exp_manager
+ logger: false # provided by exp_manager
+ accumulate_grad_batches: 1 # accumulates grads every k batches
+ gradient_clip_val: 0.0
+ amp_level: O0 # O1/O2 for mixed precision
+ precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP.
+ accelerator: ddp
+
+decoder_model:
+ do_training: true
+ transformer: t5-base
+ tokenizer: ${decoder_model.transformer}
+ nemo_path: ${decoder_exp_manager.exp_dir}/decoder_model.nemo # exported .nemo path
+
+ optim:
+ name: adamw
+ lr: 2e-4
+ weight_decay: 0.01
+
+ sched:
+ name: WarmupAnnealing
+
+ # pytorch lightning args
+ monitor: val_loss
+ reduce_on_plateau: false
+
+ # scheduler config override
+ warmup_steps: null
+ warmup_ratio: 0.0
+ last_epoch: -1
+
+decoder_exp_manager:
+ exp_dir: exps # where to store logs and checkpoints
+ name: decoder_training # name of experiment
+ create_tensorboard_logger: True
+ create_checkpoint_callback: True
+ checkpoint_callback_params:
+ save_top_k: 3
+ monitor: "val_loss"
+ mode: "min"
+ save_best_model: true
+ always_save_nemo: true
+
+# Data
+data:
+ base_dir: ??? # /path/to/data
+
+ train_ds:
+ data_path: ${data.base_dir}/train.tsv
+ batch_size: 64
+ shuffle: true
+ do_basic_tokenize: false
+ max_decoder_len: 80
+ mode: ${mode}
+ # Refer to the text_normalization doc for more information about data augmentation
+ tagger_data_augmentation: true
+ decoder_data_augmentation: true
+
+ validation_ds:
+ data_path: ${data.base_dir}/dev.tsv
+ batch_size: 64
+ shuffle: false
+ do_basic_tokenize: false
+ max_decoder_len: 80
+ mode: ${mode}
+
+ test_ds:
+ data_path: ${data.base_dir}/test.tsv
+ batch_size: 64
+ shuffle: false
+ mode: ${mode}
+
+# Inference
+inference:
+ interactive: false # Set to true if you want to enable the interactive mode when running duplex_text_normalization_test.py
+ errors_log_fp: errors.txt # Path to the file for logging the errors
diff --git a/examples/nlp/duplex_text_normalization/duplex_text_normalization_test.py b/examples/nlp/duplex_text_normalization/duplex_text_normalization_test.py
new file mode 100644
index 000000000000..9c461b18ddbb
--- /dev/null
+++ b/examples/nlp/duplex_text_normalization/duplex_text_normalization_test.py
@@ -0,0 +1,89 @@
+# Copyright (c) 2021, 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 evaluate a DuplexTextNormalizationModel.
+Note that DuplexTextNormalizationModel is essentially a wrapper class around
+DuplexTaggerModel and DuplexDecoderModel. Therefore, two trained NeMo models
+should be specificied before evaluation (one is a trained DuplexTaggerModel
+and the other is a trained DuplexDecoderModel).
+
+USAGE Example:
+1. Obtain a processed test data file (refer to the `text_normalization doc `)
+2.
+# python duplex_text_normalization_test.py
+ tagger_pretrained_model=PATH_TO_TRAINED_TAGGER
+ decoder_pretrained_model=PATH_TO_TRAINED_DECODER
+ data.test_ds.data_path=PATH_TO_TEST_FILE
+ mode={tn,itn,joint}
+
+The script also supports the `interactive` mode where a user can just make the model
+run on any input text:
+# python duplex_text_normalization_test.py
+ tagger_pretrained_model=PATH_TO_TRAINED_TAGGER
+ decoder_pretrained_model=PATH_TO_TRAINED_DECODER
+ mode={tn,itn,joint}
+ inference.interactive=true
+
+This script uses the `/examples/nlp/duplex_text_normalization/conf/duplex_tn_config.yaml`
+config file by default. The other option is to set another config file via command
+line arguments by `--config-name=CONFIG_FILE_PATH'.
+
+Note that when evaluating a DuplexTextNormalizationModel on a labeled dataset,
+the script will automatically generate a file for logging the errors made
+by the model. The location of this file is determined by the argument
+`inference.errors_log_fp`.
+
+"""
+
+
+from helpers import DECODER_MODEL, TAGGER_MODEL, instantiate_model_and_trainer
+from nltk import word_tokenize
+from omegaconf import DictConfig, OmegaConf
+
+import nemo.collections.nlp.data.text_normalization.constants as constants
+from nemo.collections.nlp.data.text_normalization import TextNormalizationTestDataset
+from nemo.collections.nlp.models import DuplexTextNormalizationModel
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+
+
+@hydra_runner(config_path="conf", config_name="duplex_tn_config")
+def main(cfg: DictConfig) -> None:
+ logging.info(f'Config Params: {OmegaConf.to_yaml(cfg)}')
+ tagger_trainer, tagger_model = instantiate_model_and_trainer(cfg, TAGGER_MODEL, False)
+ decoder_trainer, decoder_model = instantiate_model_and_trainer(cfg, DECODER_MODEL, False)
+ tn_model = DuplexTextNormalizationModel(tagger_model, decoder_model)
+
+ if not cfg.inference.interactive:
+ # Setup test_dataset
+ test_dataset = TextNormalizationTestDataset(cfg.data.test_ds.data_path, cfg.data.test_ds.mode)
+ results = tn_model.evaluate(test_dataset, cfg.data.test_ds.batch_size, cfg.inference.errors_log_fp)
+ print(f'\nTest results: {results}')
+ else:
+ while True:
+ test_input = input('Input a test input:')
+ test_input = ' '.join(word_tokenize(test_input))
+ outputs = tn_model._infer([test_input, test_input], [constants.INST_BACKWARD, constants.INST_FORWARD])[-1]
+ print(f'Prediction (ITN): {outputs[0]}')
+ print(f'Prediction (TN): {outputs[1]}')
+
+ should_continue = input('\nContinue (y/n): ').strip().lower()
+ if should_continue.startswith('n'):
+ break
+
+
+if __name__ == '__main__':
+ main()
diff --git a/examples/nlp/duplex_text_normalization/duplex_text_normalization_train.py b/examples/nlp/duplex_text_normalization/duplex_text_normalization_train.py
new file mode 100644
index 000000000000..b96451569603
--- /dev/null
+++ b/examples/nlp/duplex_text_normalization/duplex_text_normalization_train.py
@@ -0,0 +1,120 @@
+# Copyright (c) 2021, 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 a DuplexTextNormalizationModel.
+Note that DuplexTextNormalizationModel is essentially a wrapper class around
+two other classes:
+
+(1) DuplexTaggerModel is a model for identifying spans in the input that need to
+be normalized. Usually, such spans belong to semiotic classes (e.g., DATE, NUMBERS, ...).
+
+(2) DuplexDecoderModel is a model for normalizing the spans identified by the tagger.
+For example, in the text normalization (TN) problem, each span will be converted to its
+spoken form. In the inverse text normalization (ITN) problem, each span will be converted
+to its written form.
+
+Therefore, this script consists of two parts, one is for training the tagger model
+and the other is for training the decoder.
+
+This script uses the `/examples/nlp/duplex_text_normalization/conf/duplex_tn_config.yaml`
+config file by default. The other option is to set another config file via command
+line arguments by `--config-name=CONFIG_FILE_PATH'. Probably it is worth looking
+at the example config file to see the list of parameters used for training.
+
+USAGE Example:
+1. Obtain a processed dataset (refer to the `text_normalization doc `)
+2.
+# python duplex_text_normalization_train.py
+ data.base_dir=PATH_TO_DATASET_DIR
+ mode={tn,itn,joint}
+
+There are 3 different modes. `tn` mode is for training a system for TN only.
+`itn` mode is for training a system for ITN. `joint` is for training a system
+that can do both TN and ITN at the same time. Note that the above command will
+first train a tagger and then train a decoder sequentially.
+
+You can also train only a tagger (without training a decoder) by running the
+following command:
+# python duplex_text_normalization_train.py
+ data.base_dir=PATH_TO_DATASET_DIR
+ mode={tn,itn,joint}
+ decoder_model.do_training=false
+
+Or you can also train only a decoder (without training a tagger):
+# python duplex_text_normalization_train.py
+ data.base_dir=PATH_TO_DATASET_DIR
+ mode={tn,itn,joint}
+ tagger_model.do_training=false
+
+Information on the arguments:
+
+Most arguments in the example config file are quite self-explanatory (e.g.,
+`decoder_model.optim.lr` refers to the learning rate for training the decoder).
+Some arguments we want to mention are:
+
++ data.base_dir: The path to the dataset directory. It is expected that the
+directory contains three files: train.tsv, dev.tsv, and test.tsv.
+
++ tagger_model.nemo_path: This is the path where the final trained tagger model
+will be saved to.
+
++ decoder_model.nemo_path: This is the path where the final trained decoder model
+will be saved to.
+"""
+
+
+from helpers import DECODER_MODEL, TAGGER_MODEL, instantiate_model_and_trainer
+from omegaconf import DictConfig, OmegaConf
+
+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="duplex_tn_config")
+def main(cfg: DictConfig) -> None:
+ logging.info(f'Config Params: {OmegaConf.to_yaml(cfg)}')
+
+ # Train the tagger
+ if cfg.tagger_model.do_training:
+ logging.info(
+ "================================================================================================"
+ )
+ logging.info('Starting training tagger...')
+ tagger_trainer, tagger_model = instantiate_model_and_trainer(cfg, TAGGER_MODEL, True)
+ exp_manager(tagger_trainer, cfg.get('tagger_exp_manager', None))
+ tagger_trainer.fit(tagger_model)
+ if cfg.tagger_model.nemo_path:
+ tagger_model.to(tagger_trainer.accelerator.root_device)
+ tagger_model.save_to(cfg.tagger_model.nemo_path)
+ logging.info('Training finished!')
+
+ # Train the decoder
+ if cfg.decoder_model.do_training:
+ logging.info(
+ "================================================================================================"
+ )
+ logging.info('Starting training decoder...')
+ decoder_trainer, decoder_model = instantiate_model_and_trainer(cfg, DECODER_MODEL, True)
+ exp_manager(decoder_trainer, cfg.get('decoder_exp_manager', None))
+ decoder_trainer.fit(decoder_model)
+ if cfg.decoder_model.nemo_path:
+ decoder_model.to(decoder_trainer.accelerator.root_device)
+ decoder_model.save_to(cfg.decoder_model.nemo_path)
+ logging.info('Training finished!')
+
+
+if __name__ == '__main__':
+ main()
diff --git a/examples/nlp/duplex_text_normalization/google_data_preprocessing.py b/examples/nlp/duplex_text_normalization/google_data_preprocessing.py
new file mode 100644
index 000000000000..656d9cba0c26
--- /dev/null
+++ b/examples/nlp/duplex_text_normalization/google_data_preprocessing.py
@@ -0,0 +1,182 @@
+# Copyright (c) 2021, 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 can be used to process the raw data files of the Google Text Normalization dataset
+to obtain data files of the format mentioned in the `text_normalization doc `.
+Note that the script also does some preprocessing on the spoken forms of the URLs. For example,
+given the URL "Zimbio.com", the original expected spoken form in the Google dataset is
+"z_letter i_letter m_letter b_letter i_letter o_letter dot c_letter o_letter m_letter".
+However, our script will return a more concise output which is "zim bio dot com".
+
+
+USAGE Example:
+1. Download the Google TN dataset from https://www.kaggle.com/google-nlu/text-normalization
+2. Unzip the English subset (e.g., by running `tar zxvf en_with_types.tgz`). Then there will a folder named `en_with_types`.
+3. Run this script
+# python google_data_preprocessing.py \
+ --data_dir=en_with_types/ \
+ --output_dir=preprocessed/ \
+ --lang=en
+
+In this example, the final preprocessed files will be stored in the `preprocessed` folder.
+The folder should contain three files `train.tsv`, 'dev.tsv', and `test.tsv`.
+"""
+
+from argparse import ArgumentParser
+from os import listdir, mkdir
+from os.path import isdir, isfile, join
+
+import wordninja
+from helpers import flatten
+from nltk import word_tokenize
+from tqdm import tqdm
+
+import nemo.collections.nlp.data.text_normalization.constants as constants
+
+# Local Constants
+ENGLISH = 'en'
+SUPPORTED_LANGS = [ENGLISH]
+TRAIN, DEV, TEST = 'train', 'dev', 'test'
+SPLIT_NAMES = [TRAIN, DEV, TEST]
+MAX_DEV_SIZE = 25000
+
+# Helper Functions
+def read_google_data(data_dir, lang):
+ """
+ The function can be used to read the raw data files of the Google Text Normalization
+ dataset (which can be downloaded from https://www.kaggle.com/google-nlu/text-normalization)
+
+ Args:
+ data_dir: Path to the data directory. The directory should contain files of the form output-xxxxx-of-00100
+ lang: Selected language. Currently the only supported language is English.
+ Return:
+ train: A list of examples in the training set.
+ dev: A list of examples in the dev set
+ test: A list of examples in the test set
+ """
+ train, dev, test = [], [], []
+ for fn in listdir(data_dir):
+ fp = join(data_dir, fn)
+ if not isfile(fp):
+ continue
+ if not fn.startswith('output'):
+ continue
+ with open(fp, 'r', encoding='utf-8') as f:
+ # Determine the current split
+ split_nb = int(fn.split('-')[1])
+ if split_nb == 0:
+ cur_split = train
+ elif split_nb == 90:
+ cur_split = dev
+ elif split_nb == 99:
+ cur_split = test
+ else:
+ continue
+ # Loop through each line of the file
+ cur_classes, cur_tokens, cur_outputs = [], [], []
+ for linectx, line in tqdm(enumerate(f)):
+ es = line.strip().split('\t')
+ if split_nb == 99 and linectx == 100002:
+ break
+ if len(es) == 2 and es[0] == '':
+ # Update cur_split
+ cur_outputs = process_url(cur_tokens, cur_outputs, lang)
+ cur_split.append((cur_classes, cur_tokens, cur_outputs))
+ # Reset
+ cur_classes, cur_tokens, cur_outputs = [], [], []
+ continue
+ # Update the current example
+ assert len(es) == 3
+ cur_classes.append(es[0])
+ cur_tokens.append(es[1])
+ cur_outputs.append(es[2])
+ dev = dev[:MAX_DEV_SIZE]
+ train_sz, dev_sz, test_sz = len(train), len(dev), len(test)
+ print(f'train_sz: {train_sz} | dev_sz: {dev_sz} | test_sz: {test_sz}')
+ return train, dev, test
+
+
+def process_url(tokens, outputs, lang):
+ """
+ The function is used to process the spoken form of every URL in an example
+
+ Args:
+ tokens: The tokens of the written form
+ outputs: The expected outputs for the spoken form
+ lang: Selected language. Currently the only supported language is English.
+ Return:
+ outputs: The outputs for the spoken form with preprocessed URLs.
+ """
+ if lang == ENGLISH:
+ for i in range(len(tokens)):
+ t, o = tokens[i], outputs[i]
+ if o != constants.SIL_WORD and '_letter' in o:
+ o_tokens = o.split(' ')
+ all_spans, cur_span = [], []
+ for j in range(len(o_tokens)):
+ if len(o_tokens[j]) == 0:
+ continue
+ if o_tokens[j] == '_letter':
+ all_spans.append(cur_span)
+ all_spans.append([' '])
+ cur_span = []
+ else:
+ o_tokens[j] = o_tokens[j].replace('_letter', '')
+ cur_span.append(o_tokens[j])
+ if len(cur_span) > 0:
+ all_spans.append(cur_span)
+ o_tokens = flatten(all_spans)
+
+ o = ''
+ for o_token in o_tokens:
+ if len(o_token) > 1:
+ o += ' ' + o_token + ' '
+ else:
+ o += o_token
+ o = o.strip()
+ o_tokens = wordninja.split(o)
+ o = ' '.join(o_tokens)
+
+ outputs[i] = o
+
+ return outputs
+
+
+# Main code
+if __name__ == '__main__':
+ parser = ArgumentParser(description='Preprocess Google text normalization dataset')
+ parser.add_argument('--data_dir', type=str, required=True, help='Path to folder with data')
+ parser.add_argument('--output_dir', type=str, default='preprocessed', help='Path to folder with preprocessed data')
+ parser.add_argument('--lang', type=str, default=ENGLISH, choices=SUPPORTED_LANGS, help='Language')
+ args = parser.parse_args()
+
+ # Create the output dir (if not exist)
+ if not isdir(args.output_dir):
+ mkdir(args.output_dir)
+
+ # Processing
+ train, dev, test = read_google_data(args.data_dir, args.lang)
+ for split, data in zip(SPLIT_NAMES, [train, dev, test]):
+ output_f = open(join(args.output_dir, f'{split}.tsv'), 'w+', encoding='utf-8')
+ for inst in data:
+ cur_classes, cur_tokens, cur_outputs = inst
+ for c, t, o in zip(cur_classes, cur_tokens, cur_outputs):
+ t = ' '.join(word_tokenize(t))
+ if not o in constants.SPECIAL_WORDS:
+ o_tokens = word_tokenize(o)
+ o_tokens = [o_tok for o_tok in o_tokens if o_tok != constants.SIL_WORD]
+ o = ' '.join(o_tokens)
+ output_f.write(f'{c}\t{t}\t{o}\n')
+ output_f.write('\t\n')
diff --git a/examples/nlp/duplex_text_normalization/helpers.py b/examples/nlp/duplex_text_normalization/helpers.py
new file mode 100644
index 000000000000..68be578f3dae
--- /dev/null
+++ b/examples/nlp/duplex_text_normalization/helpers.py
@@ -0,0 +1,73 @@
+# Copyright (c) 2021, 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 pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.nlp.models import DuplexDecoderModel, DuplexTaggerModel
+from nemo.utils import logging
+
+__all__ = ['TAGGER_MODEL', 'DECODER_MODEL', 'MODEL_NAMES', 'instantiate_model_and_trainer']
+
+TAGGER_MODEL = 'tagger'
+DECODER_MODEL = 'decoder'
+MODEL_NAMES = [TAGGER_MODEL, DECODER_MODEL]
+
+
+def instantiate_model_and_trainer(cfg: DictConfig, model_name: str, do_training: bool):
+ """ Function for instantiating a model and a trainer
+ Args:
+ cfg: The config used to instantiate the model and the trainer.
+ model_name: A str indicates whether the model to be instantiated is a tagger or a decoder (i.e., model_name should be either TAGGER_MODEL or DECODER_MODEL).
+ do_training: A boolean flag indicates whether the model will be trained or evaluated.
+
+ Returns:
+ trainer: A PyTorch Lightning trainer
+ model: A NLPModel that can either be a DuplexTaggerModel or a DuplexDecoderModel
+ """
+ assert model_name in MODEL_NAMES
+ logging.info(f'Model {model_name}')
+
+ # Get configs for the corresponding models
+ trainer_cfg = cfg.get(f'{model_name}_trainer')
+ model_cfg = cfg.get(f'{model_name}_model')
+ pretrained_cfg = cfg.get(f'{model_name}_pretrained_model', None)
+
+ trainer = pl.Trainer(**trainer_cfg)
+
+ if not pretrained_cfg:
+ logging.info(f'Config: {OmegaConf.to_yaml(cfg)}')
+ if model_name == TAGGER_MODEL:
+ model = DuplexTaggerModel(model_cfg, trainer=trainer)
+ if model_name == DECODER_MODEL:
+ model = DuplexDecoderModel(model_cfg, trainer=trainer)
+ else:
+ logging.info(f'Loading pretrained model {pretrained_cfg}')
+ if model_name == TAGGER_MODEL:
+ model = DuplexTaggerModel.restore_from(pretrained_cfg)
+ if model_name == DECODER_MODEL:
+ model = DuplexDecoderModel.restore_from(pretrained_cfg)
+
+ # Setup train and validation data
+ if do_training:
+ model.setup_training_data(train_data_config=cfg.data.train_ds)
+ model.setup_validation_data(val_data_config=cfg.data.validation_ds)
+
+ logging.info(f'Model Device {model.device}')
+ return trainer, model
+
+
+def flatten(l):
+ """ flatten a list of lists """
+ return [item for sublist in l for item in sublist]
diff --git a/nemo/collections/nlp/data/text_normalization/__init__.py b/nemo/collections/nlp/data/text_normalization/__init__.py
new file mode 100644
index 000000000000..8cf835fbe0c1
--- /dev/null
+++ b/nemo/collections/nlp/data/text_normalization/__init__.py
@@ -0,0 +1,16 @@
+# Copyright (c) 2021, 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 nemo.collections.nlp.data.text_normalization.decoder_dataset import TextNormalizationDecoderDataset
+from nemo.collections.nlp.data.text_normalization.tagger_dataset import TextNormalizationTaggerDataset
+from nemo.collections.nlp.data.text_normalization.test_dataset import TextNormalizationTestDataset
diff --git a/nemo/collections/nlp/data/text_normalization/constants.py b/nemo/collections/nlp/data/text_normalization/constants.py
new file mode 100644
index 000000000000..5c6c572c7736
--- /dev/null
+++ b/nemo/collections/nlp/data/text_normalization/constants.py
@@ -0,0 +1,103 @@
+# Copyright (c) 2021, 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.
+
+DECODE_CTX_SIZE = 3 # the size of the input context to be provided to the DuplexDecoderModel
+LABEL_PAD_TOKEN_ID = -100
+
+# Task Prefixes
+ITN_PREFIX = str(0)
+TN_PREFIX = str(1)
+
+# Tagger Labels Prefixes
+B_PREFIX = 'B-' # Denote beginning
+I_PREFIX = 'I-' # Denote middle
+TAGGER_LABELS_PREFIXES = [B_PREFIX, I_PREFIX]
+
+# Modes
+TN_MODE = 'tn'
+ITN_MODE = 'itn'
+JOINT_MODE = 'joint'
+MODES = [TN_MODE, ITN_MODE, JOINT_MODE]
+
+# Instance Directions
+INST_BACKWARD = 'BACKWARD'
+INST_FORWARD = 'FORWARD'
+INST_DIRECTIONS = [INST_BACKWARD, INST_FORWARD]
+
+# TAGS
+SAME_TAG = 'SAME' # Tag indicates that a token can be kept the same without any further transformation
+TASK_TAG = 'TASK' # Tag indicates that a token belongs to a task prefix (the prefix indicates whether the current task is TN or ITN)
+PUNCT_TAG = 'PUNCT' # Tag indicates that a token is a punctuation
+TRANSFORM_TAG = 'TRANSFORM' # Tag indicates that a token needs to be transformed by the decoder
+ALL_TAGS = [TASK_TAG, SAME_TAG, PUNCT_TAG, TRANSFORM_TAG]
+
+# ALL_TAG_LABELS
+ALL_TAG_LABELS = []
+for prefix in TAGGER_LABELS_PREFIXES:
+ for tag in ALL_TAGS:
+ ALL_TAG_LABELS.append(prefix + tag)
+ALL_TAG_LABELS.sort()
+
+# Special Words
+SIL_WORD = 'sil'
+SELF_WORD = ''
+SPECIAL_WORDS = [SIL_WORD, SELF_WORD]
+
+# Mappings for Greek Letters
+GREEK_TO_SPOKEN = {
+ 'Τ': 'tau',
+ 'Ο': 'omicron',
+ 'Δ': 'delta',
+ 'Η': 'eta',
+ 'Κ': 'kappa',
+ 'Ι': 'iota',
+ 'Θ': 'theta',
+ 'Α': 'alpha',
+ 'Σ': 'sigma',
+ 'Υ': 'upsilon',
+ 'Μ': 'mu',
+ 'Ε': 'epsilon',
+ 'Χ': 'chi',
+ 'Π': 'pi',
+ 'Ν': 'nu',
+ 'Λ': 'lambda',
+ 'Γ': 'gamma',
+ 'Β': 'beta',
+ 'Ρ': 'rho',
+ 'τ': 'tau',
+ 'υ': 'upsilon',
+ 'μ': 'mu',
+ 'φ': 'phi',
+ 'α': 'alpha',
+ 'λ': 'lambda',
+ 'ι': 'iota',
+ 'ς': 'sigma',
+ 'ο': 'omicron',
+ 'σ': 'sigma',
+ 'η': 'eta',
+ 'π': 'pi',
+ 'ν': 'nu',
+ 'γ': 'gamma',
+ 'κ': 'kappa',
+ 'ε': 'epsilon',
+ 'β': 'beta',
+ 'ρ': 'rho',
+ 'ω': 'omega',
+ 'χ': 'chi',
+}
+SPOKEN_TO_GREEK = {v: k for k, v in GREEK_TO_SPOKEN.items()}
+
+# IDs for special tokens for encoding inputs of the decoder models
+EXTRA_ID_0 = ''
+EXTRA_ID_1 = ''
diff --git a/nemo/collections/nlp/data/text_normalization/decoder_dataset.py b/nemo/collections/nlp/data/text_normalization/decoder_dataset.py
new file mode 100644
index 000000000000..30f3b5799d6d
--- /dev/null
+++ b/nemo/collections/nlp/data/text_normalization/decoder_dataset.py
@@ -0,0 +1,196 @@
+# Copyright (c) 2021, 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 random
+
+from nltk import word_tokenize
+from tqdm import tqdm
+from transformers import PreTrainedTokenizerBase
+
+import nemo.collections.nlp.data.text_normalization.constants as constants
+from nemo.collections.nlp.data.text_normalization.utils import read_data_file
+from nemo.core.classes import Dataset
+from nemo.utils.decorators.experimental import experimental
+
+__all__ = ['TextNormalizationDecoderDataset']
+
+
+@experimental
+class TextNormalizationDecoderDataset(Dataset):
+ """
+ Creates dataset to use to train a DuplexDecoderModel.
+
+ Converts from raw data to an instance that can be used by Dataloader.
+
+ For dataset to use to do end-to-end inference, see TextNormalizationTestDataset.
+
+ Args:
+ input_file: path to the raw data file (e.g., train.tsv). For more info about the data format, refer to the `text_normalization doc `.
+ tokenizer: tokenizer of the model that will be trained on the dataset
+ mode: should be one of the values ['tn', 'itn', 'joint']. `tn` mode is for TN only. `itn` mode is for ITN only. `joint` is for training a system that can do both TN and ITN at the same time.
+ max_len: maximum length of sequence in tokens. The code will discard any training instance whose input or output is longer than the specified max_len.
+ decoder_data_augmentation (bool): a flag indicates whether to augment the dataset with additional data instances that may help the decoder become more robust against the tagger's errors. Refer to the doc for more info.
+ """
+
+ def __init__(
+ self,
+ input_file: str,
+ tokenizer: PreTrainedTokenizerBase,
+ mode: str,
+ max_len: int,
+ decoder_data_augmentation: bool,
+ ):
+ assert mode in constants.MODES
+ self.mode = mode
+ raw_insts = read_data_file(input_file)
+
+ # Convert raw instances to TaggerDataInstance
+ insts, inputs, targets = [], [], []
+ for (classes, w_words, s_words) in tqdm(raw_insts):
+ for ix, (_class, w_word, s_word) in enumerate(zip(classes, w_words, s_words)):
+ if s_word in constants.SPECIAL_WORDS:
+ continue
+ for inst_dir in constants.INST_DIRECTIONS:
+ if inst_dir == constants.INST_BACKWARD and mode == constants.TN_MODE:
+ continue
+ if inst_dir == constants.INST_FORWARD and mode == constants.ITN_MODE:
+ continue
+ # Create a DecoderDataInstance
+ inst = DecoderDataInstance(
+ w_words, s_words, inst_dir, start_idx=ix, end_idx=ix + 1, semiotic_class=_class
+ )
+ insts.append(inst)
+ if decoder_data_augmentation:
+ noise_left = random.randint(1, 2)
+ noise_right = random.randint(1, 2)
+ inst = DecoderDataInstance(
+ w_words, s_words, inst_dir, start_idx=ix - noise_left, end_idx=ix + 1 + noise_right
+ )
+ insts.append(inst)
+
+ self.insts = insts
+ inputs = [inst.input_str for inst in insts]
+ targets = [inst.output_str for inst in insts]
+
+ # Tokenization
+ self.inputs, self.examples = [], []
+ self.tn_count, self.itn_count, long_examples_filtered = 0, 0, 0
+ input_max_len, target_max_len = 0, 0
+ for idx in range(len(inputs)):
+ # Input
+ _input = tokenizer([inputs[idx]])
+ input_len = len(_input['input_ids'][0])
+ if input_len > max_len:
+ long_examples_filtered += 1
+ continue
+
+ # Target
+ _target = tokenizer([targets[idx]])
+ target_len = len(_target['input_ids'][0])
+ if target_len > max_len:
+ long_examples_filtered += 1
+ continue
+
+ # Update
+ self.inputs.append(inputs[idx])
+ _input['labels'] = _target['input_ids']
+ self.examples.append(_input)
+ if inputs[idx].startswith(constants.TN_PREFIX):
+ self.tn_count += 1
+ if inputs[idx].startswith(constants.ITN_PREFIX):
+ self.itn_count += 1
+ input_max_len = max(input_max_len, input_len)
+ target_max_len = max(target_max_len, target_len)
+ print(f'long_examples_filtered: {long_examples_filtered}')
+ print(f'input_max_len: {input_max_len} | target_max_len: {target_max_len}')
+
+ def __getitem__(self, idx):
+ example = self.examples[idx]
+ item = {key: val[0] for key, val in example.items()}
+ return item
+
+ def __len__(self):
+ return len(self.examples)
+
+
+class DecoderDataInstance:
+ """
+ This class represents a data instance in a TextNormalizationDecoderDataset.
+
+ Intuitively, each data instance can be thought as having the following form:
+ Input:
+ Output: