adding guard (#3655)

Signed-off-by: Yang Zhang <[email protected]>
NVIDIA · Feb 14, 2022 · fbbfb08 · fbbfb08
1 parent 6a517f0
commit fbbfb08
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 2 deletions.
diff --git a/nemo_text_processing/inverse_text_normalization/inverse_normalize.py b/nemo_text_processing/inverse_text_normalization/inverse_normalize.py
@@ -16,6 +16,7 @@
 from time import perf_counter
 from typing import List
 
+from nemo_text_processing.text_normalization.data_loader_utils import check_installation, get_installation_msg
 from nemo_text_processing.text_normalization.normalize import Normalizer
 from nemo_text_processing.text_normalization.token_parser import TokenParser
 
@@ -32,6 +33,9 @@ class InverseNormalizer(Normalizer):
     """
 
     def __init__(self, lang: str = 'en', cache_dir: str = None, overwrite_cache: bool = False):
+
+        if not check_installation():
+            raise ImportError(get_installation_msg())
         if lang == 'en':
             from nemo_text_processing.inverse_text_normalization.en.taggers.tokenize_and_classify import ClassifyFst
             from nemo_text_processing.inverse_text_normalization.en.verbalizers.verbalize_final import (

diff --git a/nemo_text_processing/text_normalization/data_loader_utils.py b/nemo_text_processing/text_normalization/data_loader_utils.py
@@ -238,3 +238,19 @@ def pre_process(text: str) -> str:
     # remove extra space
     text = re.sub(r' +', ' ', text)
     return text
+
+
+def check_installation():
+    try:
+        import pynini
+
+        PYNINI_AVAILABLE = True
+
+    except (ModuleNotFoundError, ImportError):
+        PYNINI_AVAILABLE = False
+    return PYNINI_AVAILABLE
+
+
+def get_installation_msg():
+    msg = "`pynini` is not installed ! \n Please run the `nemo_text_processing/setup.sh` script prior to usage of this toolkit."
+    return msg
diff --git a/nemo_text_processing/text_normalization/en/graph_utils.py b/nemo_text_processing/text_normalization/en/graph_utils.py
@@ -116,7 +116,7 @@
     PYNINI_AVAILABLE = False
 
 
-def generator_main(file_name: str, graphs: Dict[str, pynini.FstLike]):
+def generator_main(file_name: str, graphs: Dict[str, 'pynini.FstLike']):
     """
     Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name.
 

diff --git a/nemo_text_processing/text_normalization/normalize.py b/nemo_text_processing/text_normalization/normalize.py
@@ -20,7 +20,7 @@
 from math import factorial
 from typing import Dict, List, Union
 
-from nemo_text_processing.text_normalization.data_loader_utils import pre_process
+from nemo_text_processing.text_normalization.data_loader_utils import get_installation_msg, pre_process
 from nemo_text_processing.text_normalization.token_parser import PRESERVE_ORDER_KEY, TokenParser
 from tqdm import tqdm
 
@@ -68,6 +68,9 @@ def __init__(
     ):
         assert input_case in ["lower_cased", "cased"]
 
+        if not PYNINI_AVAILABLE:
+            raise ImportError(get_installation_msg())
+
         if lang == 'en' and deterministic:
             from nemo_text_processing.text_normalization.en.taggers.tokenize_and_classify import ClassifyFst
             from nemo_text_processing.text_normalization.en.verbalizers.verbalize_final import VerbalizeFinalFst