update english tn ckpt (NVIDIA#3143)

* update english tn ckpt Signed-off-by: Yang Zhang <[email protected]> * remove ununsed import Signed-off-by: Yang Zhang <[email protected]>
ahmetgunduz · Nov 5, 2021 · 875f544 · 875f544
1 parent a6834f8
commit 875f544
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 4 deletions.
diff --git a/docs/source/nlp/text_normalization.rst b/docs/source/nlp/text_normalization.rst
@@ -204,9 +204,9 @@ Run the full inference pipeline:
     cd NeMo/examples/nlp/duplex_text_normalization;
 
     # run inference in interactive mode using pretrained tagger and decoder models
-    python duplex_text_normalization duplex_text_normalization_infer.py \
-        tagger_pretrained_model=text_normalization_tagger_en \
-        decoder_pretrained_model=text_normalization_decoder_en \
+    python duplex_text_normalization_infer.py \
+        tagger_pretrained_model=neural_text_normalization_t5 \
+        decoder_pretrained_model=neural_text_normalization_t5 \
         inference.from_file=False \
         lang=en \
         mode=tn

diff --git a/nemo/collections/nlp/models/duplex_text_normalization/duplex_decoder.py b/nemo/collections/nlp/models/duplex_text_normalization/duplex_decoder.py
@@ -526,4 +526,11 @@ def list_available_models(cls) -> Optional[PretrainedModelInfo]:
             List of available pre-trained models.
         """
         result = []
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="neural_text_normalization_t5",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/neural_text_normalization_t5/versions/1.5.0/files/neural_text_normalization_t5_decoder.nemo",
+                description="Text Normalization model's decoder model.",
+            )
+        )
         return result
diff --git a/nemo/collections/nlp/models/duplex_text_normalization/duplex_tagger.py b/nemo/collections/nlp/models/duplex_text_normalization/duplex_tagger.py
@@ -382,4 +382,11 @@ def list_available_models(cls) -> Optional[PretrainedModelInfo]:
             List of available pre-trained models.
         """
         result = []
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="neural_text_normalization_t5",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/neural_text_normalization_t5/versions/1.5.0/files/neural_text_normalization_t5_tagger.nemo",
+                description="Text Normalization model's tagger model.",
+            )
+        )
         return result
diff --git a/nemo/collections/nlp/models/duplex_text_normalization/duplex_tn.py b/nemo/collections/nlp/models/duplex_text_normalization/duplex_tn.py
@@ -24,7 +24,6 @@
 from nemo.collections.nlp.data.text_normalization.utils import input_preprocessing, post_process_punct
 from nemo.collections.nlp.models.duplex_text_normalization.utils import get_formatted_string
 from nemo.utils import logging
-from nemo.utils.decorators.experimental import experimental
 
 __all__ = ['DuplexTextNormalizationModel', 'post_process_punct']