fix(tn/itn): add cache_dir for cli (#131)

wenet-e2e · Oct 30, 2023 · be9aa9c · be9aa9c
1 parent 94c3eb2
commit be9aa9c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/itn/main.py b/itn/main.py
@@ -30,6 +30,9 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--text', help='input string')
     parser.add_argument('--file', help='input file path')
+    parser.add_argument('--cache_dir', type=str,
+                        default=None,
+                        help='cache dir containing *.fst')
     parser.add_argument('--overwrite_cache', action='store_true',
                         help='rebuild *.fst')
     parser.add_argument('--enable_standalone_number', type=str,
@@ -41,7 +44,7 @@ def main():
     args = parser.parse_args()
 
     normalizer = InverseNormalizer(
-        cache_dir='itn', overwrite_cache=args.overwrite_cache,
+        cache_dir=args.cache_dir, overwrite_cache=args.overwrite_cache,
         enable_standalone_number=str2bool(args.enable_standalone_number),
         enable_0_to_9=str2bool(args.enable_0_to_9))
 

diff --git a/tn/main.py b/tn/main.py
@@ -22,11 +22,14 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--text', help='input string')
     parser.add_argument('--file', help='input file path')
+    parser.add_argument('--cache_dir', type=str,
+                        default=None,
+                        help='cache dir containing *.fst')
     parser.add_argument('--overwrite_cache', action='store_true',
                         help='rebuild *.fst')
     args = parser.parse_args()
 
-    normalizer = Normalizer(cache_dir='tn',
+    normalizer = Normalizer(cache_dir=args.cache_dir,
                             overwrite_cache=args.overwrite_cache)
 
     if args.text: