refactor: harden TN and ITN processing

Author: pengzhendong

.github/workflows/wheels.yml

@@ -32,12 +32,15 @@ jobs:
           python -m tn --text "2.5平方电线" --overwrite_cache --language "zh"
           python -m tn --text "2010-03-21" --overwrite_cache --language "en"
           python -m itn --text "二点五平方电线" --overwrite_cache
+          python -m itn --text "one two three" --overwrite_cache --language "en"
 
       - name: Prepare Graph
         run: |
           mkdir graph
           cp tn/*.fst graph
           cp itn/*.fst graph
+          cp tn/*_cache.json graph
+          cp itn/*_cache.json graph
 
       - name: Get version from setuptools_scm
         id: scm_version

.gitignore

@@ -33,6 +33,7 @@ __pycache__/
 
 # Clangd files
 .cache
+.venv/
 compile_commands.json
 
 

README.md

@@ -34,22 +34,26 @@ weitn --text "二点五平方电线"
 Python usage:
 
 ```py
-from itn.chinese.inverse_normalizer import InverseNormalizer
-from tn.chinese.normalizer import Normalizer as ZhNormalizer
-from tn.english.normalizer import Normalizer as EnNormalizer
-
-# NOTE(xcsong): 和默认参数不一致时，必须重新构图，要重新构图请务必指定 `overwrite_cache=True`
-#               When the parameters differ from the defaults, it is mandatory to re-compose. To re-compose, please ensure you specify `overwrite_cache=True`.
+from itn.chinese.inverse_normalizer import InverseNormalizer
+from itn.english.inverse_normalizer import InverseNormalizer as EnInverseNormalizer
+from tn.chinese.normalizer import Normalizer as ZhNormalizer
+from tn.english.normalizer import Normalizer as EnNormalizer
+
+# FST 缓存会记录构建参数及规则数据指纹；配置或规则变化时会自动重新构图。
+# Set `overwrite_cache=True` only when an unconditional rebuild is required.
 
 zh_tn_text = "你好 WeTextProcessing 1.0，船新版本儿，船新体验儿，简直666，9和10"
 zh_itn_text = "你好 WeTextProcessing 一点零，船新版本儿，船新体验儿，简直六六六，九和六"
-en_tn_text = "Hello WeTextProcessing 1.0, life is short, just use wetext, 666, 9 and 10"
+en_tn_text = "Hello WeTextProcessing 1.0, life is short, just use wetext, 666, 9 and 10"
+en_itn_text = "call me at five five five one two three four"
 zh_tn_model = ZhNormalizer(remove_erhua=True, overwrite_cache=True)
 zh_itn_model = InverseNormalizer(enable_0_to_9=False, overwrite_cache=True)
-en_tn_model = EnNormalizer(overwrite_cache=True)
+en_tn_model = EnNormalizer(overwrite_cache=True)
+en_itn_model = EnInverseNormalizer(overwrite_cache=True)
 print("中文 TN (去除儿化音，重新在线构图):\n\t{} => {}".format(zh_tn_text, zh_tn_model.normalize(zh_tn_text)))
 print("中文ITN (小于10的单独数字不转换，重新在线构图):\n\t{} => {}".format(zh_itn_text, zh_itn_model.normalize(zh_itn_text)))
-print("英文 TN (暂时还没有可控的选项，后面会加...):\n\t{} => {}\n".format(en_tn_text, en_tn_model.normalize(en_tn_text)))
+print("英文 TN (暂时还没有可控的选项，后面会加...):\n\t{} => {}\n".format(en_tn_text, en_tn_model.normalize(en_tn_text)))
+print("英文 ITN:\n\t{} => {}\n".format(en_itn_text, en_itn_model.normalize(en_itn_text)))
 
 zh_tn_model = ZhNormalizer(overwrite_cache=False)
 zh_itn_model = InverseNormalizer(overwrite_cache=False)

itn/chinese/inverse_normalizer.py

@@ -48,7 +48,17 @@ def __init__(
         self.enable_million = enable_million
         if cache_dir is None:
             cache_dir = files("itn")
-        self.build_fst("zh_itn", cache_dir, overwrite_cache)
+        self.build_fst(
+            "zh_itn",
+            cache_dir,
+            overwrite_cache,
+            {
+                "enable_0_to_9": self.enable_0_to_9,
+                "enable_million": self.enable_million,
+                "enable_standalone_number": self.convert_number,
+                "remove_interjections": self.remove_interjections,
+            },
+        )
 
     def build_tagger_and_verbalizer(self):
         cardinal = Cardinal(self.convert_number, self.enable_0_to_9, self.enable_million)

itn/english/inverse_normalizer.py

@@ -38,7 +38,7 @@ def __init__(self, cache_dir=None, overwrite_cache=False):
         super().__init__(name="en_inverse_normalizer", ordertype="itn")
         if cache_dir is None:
             cache_dir = files("itn")
-        self.build_fst("en_itn", cache_dir, overwrite_cache)
+        self.build_fst("en_itn", cache_dir, overwrite_cache, {})
 
     def build_tagger_and_verbalizer(self):
         cardinal = Cardinal()

itn/japanese/inverse_normalizer.py

@@ -48,7 +48,17 @@ def __init__(
         self.enable_million = enable_million
         if cache_dir is None:
             cache_dir = files("itn")
-        self.build_fst("ja_itn", cache_dir, overwrite_cache)
+        self.build_fst(
+            "ja_itn",
+            cache_dir,
+            overwrite_cache,
+            {
+                "enable_0_to_9": self.enable_0_to_9,
+                "enable_million": self.enable_million,
+                "enable_standalone_number": self.convert_number,
+                "full_to_half": self.full_to_half,
+            },
+        )
 
     def build_tagger_and_verbalizer(self):
         processor = PreProcessor(full_to_half=self.full_to_half).processor

itn/main.py

@@ -15,6 +15,7 @@
 import argparse
 
 from itn.chinese.inverse_normalizer import InverseNormalizer as ZhInverseNormalizer
+from itn.english.inverse_normalizer import InverseNormalizer as EnInverseNormalizer
 from itn.japanese.inverse_normalizer import InverseNormalizer as JaInverseNormalizer
 from tn.utils import str2bool
 
@@ -28,7 +29,7 @@ def main():
     parser.add_argument("--enable_standalone_number", type=str, default="True", help="一百 = 100 if True else 一百")
     parser.add_argument("--enable_0_to_9", type=str, default="False", help="零和九 = 0和9 if True else 零和九")
     parser.add_argument("--enable_million", type=str, default="False", help="六百万 = 6000000 if True else 600万")
-    parser.add_argument("--language", type=str, default="zh", choices=["zh", "ja"], help="valid languages")
+    parser.add_argument("--language", type=str, default="zh", choices=["zh", "en", "ja"], help="valid languages")
     args = parser.parse_args()
 
     if args.language == "zh":
@@ -39,6 +40,11 @@ def main():
             enable_0_to_9=str2bool(args.enable_0_to_9),
             enable_million=str2bool(args.enable_million),
         )
+    elif args.language == "en":
+        normalizer = EnInverseNormalizer(
+            cache_dir=args.cache_dir,
+            overwrite_cache=args.overwrite_cache,
+        )
     elif args.language == "ja":
         normalizer = JaInverseNormalizer(
             cache_dir=args.cache_dir,

pyproject.toml

@@ -53,5 +53,12 @@ include = ["tn*", "itn*"]
 namespaces = false
 
 [tool.setuptools.package-data]
-tn = ["*.fst", "chinese/data/*/*.tsv", "english/data/*/*.tsv", "english/data/*.tsv", "english/data/*/*.far", "japanese/data/*/*.tsv"]
-itn = ["*.fst", "chinese/data/*/*.tsv", "japanese/data/*/*.tsv"]
+tn = ["*.fst", "*_cache.json", "chinese/data/*/*.tsv", "english/data/*/*.tsv", "english/data/*.tsv", "english/data/*/*.far", "japanese/data/*/*.tsv"]
+itn = [
+    "*.fst",
+    "*_cache.json",
+    "chinese/data/*/*.tsv",
+    "english/data/*.tsv",
+    "english/data/*/*.tsv",
+    "japanese/data/*/*.tsv",
+]

runtime/processor/wetext_processor.cc

@@ -30,8 +30,12 @@ Processor::Processor(const std::string& tagger_path,
     parse_type_ = ParseType::kZH_ITN;
   } else if (tagger_path.find("en_tn_") != tagger_path.npos) {
     parse_type_ = ParseType::kEN_TN;
+  } else if (tagger_path.find("en_itn_") != tagger_path.npos) {
+    parse_type_ = ParseType::kEN_ITN;
   } else if (tagger_path.find("ja_tn_") != tagger_path.npos) {
     parse_type_ = ParseType::kJA_TN;
+  } else if (tagger_path.find("ja_itn_") != tagger_path.npos) {
+    parse_type_ = ParseType::kJA_ITN;
   } else {
     LOG(FATAL) << "Invalid fst prefix, prefix should contain"
                << " either \"_tn_\" or \"_itn_\".";

runtime/processor/wetext_token_parser.cc

@@ -14,6 +14,8 @@
 
 #include "processor/wetext_token_parser.h"
 
+#include <stdexcept>
+
 #include "utils/wetext_log.h"
 #include "utils/wetext_string.h"
 
@@ -33,24 +35,26 @@ const std::unordered_map<std::string, std::vector<std::string>> ZH_TN_ORDERS = {
     {"money", {"value", "currency"}},
     {"time", {"noon", "hour", "minute", "second"}}};
 const std::unordered_map<std::string, std::vector<std::string>> JA_TN_ORDERS = {
-    {"date", {"year", "month", "day"}},
-    {"money", {"value", "currency"}}};
+    {"date", {"year", "month", "day"}}, {"money", {"value", "currency"}}};
 
 const std::unordered_map<std::string, std::vector<std::string>> EN_TN_ORDERS = {
     {"date", {"preserve_order", "text", "day", "month", "year"}},
     {"money", {"integer_part", "fractional_part", "quantity", "currency_maj"}}};
-const std::unordered_map<std::string, std::vector<std::string>> ZH_ITN_ORDERS =
-    {{"date", {"year", "month", "day"}},
-     {"fraction", {"sign", "numerator", "denominator"}},
-     {"measure", {"numerator", "denominator", "value"}},
-     {"money", {"currency", "value", "decimal"}},
-     {"time", {"hour", "minute", "second", "noon"}}};
+const std::unordered_map<std::string, std::vector<std::string>> ITN_ORDERS = {
+    {"date", {"year", "month", "day", "preserve_order"}},
+    {"fraction", {"sign", "numerator", "denominator"}},
+    {"measure", {"numerator", "denominator", "value", "units"}},
+    {"money", {"currency", "value", "decimal", "quantity"}},
+    {"time", {"hour", "minute", "second", "noon", "zone"}},
+    {"telephone", {"country_code", "number_part"}},
+    {"electronic", {"username", "domain", "protocol"}}};
 
 TokenParser::TokenParser(ParseType type) {
   if (type == ParseType::kZH_TN) {
     orders_ = ZH_TN_ORDERS;
-  } else if (type == ParseType::kZH_ITN) {
-    orders_ = ZH_ITN_ORDERS;
+  } else if (type == ParseType::kZH_ITN || type == ParseType::kEN_ITN ||
+             type == ParseType::kJA_ITN) {
+    orders_ = ITN_ORDERS;
   } else if (type == ParseType::kEN_TN) {
     orders_ = EN_TN_ORDERS;
   } else if (type == ParseType::kJA_TN) {
@@ -62,9 +66,12 @@ TokenParser::TokenParser(ParseType type) {
 
 void TokenParser::Load(const std::string& input) {
   wetext::SplitUTF8StringToChars(input, &text_);
-  CHECK_GT(text_.size(), 0);
+  if (text_.empty()) {
+    throw std::invalid_argument("token stream must not be empty");
+  }
   index_ = 0;
   ch_ = text_[0];
+  tokens_.clear();
 }
 
 bool TokenParser::Read() {
@@ -94,38 +101,54 @@ bool TokenParser::ParseChar(const std::string& exp) {
 }
 
 bool TokenParser::ParseChars(const std::string& exp) {
-  bool ok = false;
+  size_t start = index_;
   std::vector<std::string> chars;
   wetext::SplitUTF8StringToChars(exp, &chars);
   for (const auto& x : chars) {
-    ok |= ParseChar(x);
+    if (!ParseChar(x)) {
+      index_ = start;
+      ch_ = text_[start];
+      return false;
+    }
   }
-  return ok;
+  return true;
 }
 
 std::string TokenParser::ParseKey() {
-  CHECK_NE(ch_, EOS);
-  CHECK_EQ(UTF8_WHITESPACE.count(ch_), 0);
+  if (ch_ == EOS || UTF8_WHITESPACE.count(ch_) > 0) {
+    throw std::invalid_argument("expected token key at position " +
+                                std::to_string(index_));
+  }
 
   std::string key = "";
   while (ASCII_LETTERS.count(ch_) > 0) {
     key += ch_;
     Read();
   }
+  if (key.empty()) {
+    throw std::invalid_argument("invalid token key at position " +
+                                std::to_string(index_));
+  }
   return key;
 }
 
 std::string TokenParser::ParseValue() {
-  CHECK_NE(ch_, EOS);
-  bool escape = false;
+  if (ch_ == EOS) {
+    throw std::invalid_argument("expected token value at end of stream");
+  }
 
   std::string value = "";
   while (ch_ != "\"") {
+    if (ch_ == EOS) {
+      throw std::invalid_argument("unterminated token value");
+    }
     value += ch_;
-    escape = ch_ == "\\";
+    bool escape = ch_ == "\\";
     Read();
     if (escape) {
-      escape = false;
+      if (ch_ == EOS) {
+        throw std::invalid_argument("unterminated escape in token value");
+      }
       value += ch_;
       Read();
     }
@@ -137,20 +160,31 @@ void TokenParser::Parse(const std::string& input) {
   Load(input);
   while (ParseWs()) {
     std::string name = ParseKey();
-    ParseChars(" { ");
+    if (!ParseChars(" { ")) {
+      throw std::invalid_argument("expected token opening delimiter");
+    }
 
     Token token(name);
+    bool closed = false;
     while (ParseWs()) {
       if (ch_ == "}") {
         ParseChar("}");
+        closed = true;
         break;
       }
       std::string key = ParseKey();
-      ParseChars(": \"");
+      if (!ParseChars(": \"")) {
+        throw std::invalid_argument("expected token field delimiter");
+      }
       std::string value = ParseValue();
-      ParseChar("\"");
+      if (!ParseChar("\"")) {
+        throw std::invalid_argument("expected closing quote");
+      }
       token.Append(key, value);
     }
+    if (!closed) {
+      throw std::invalid_argument("unterminated token " + name);
+    }
     tokens_.emplace_back(token);
   }
 }