📑 Add AutoProcessor Class.

Author: dathudeptrai

tensorflow_tts/inference/__init__.py

@@ -1,2 +1,3 @@
 from tensorflow_tts.inference.auto_model import TFAutoModel
 from tensorflow_tts.inference.auto_config import AutoConfig
+from tensorflow_tts.inference.auto_processor import AutoProcessor

tensorflow_tts/inference/auto_config.py

@@ -56,8 +56,8 @@ def from_pretrained(cls, pretrained_path, **kwargs):
             return config_class
         except Exception:
             raise ValueError(
-                "Unrecognized model in {}. "
-                "Should have a `model_type` key in its config.json, or contain one of the following strings "
+                "Unrecognized config in {}. "
+                "Should have a `model_type` key in its config.yaml, or contain one of the following strings "
                 "in its name: {}".format(
                     pretrained_path, ", ".join(CONFIG_MAPPING.keys())
                 )

tensorflow_tts/inference/auto_processor.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+# Copyright 2020 The TensorFlowTTS Team.
+#
+# 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.
+"""Tensorflow Auto Processor modules."""
+
+import logging
+import json
+from collections import OrderedDict
+
+from tensorflow_tts.processor import (
+    LJSpeechProcessor,
+    KSSProcessor,
+    BakerProcessor,
+    LibriTTSProcessor,
+)
+
+CONFIG_MAPPING = OrderedDict(
+    [
+        ("LJSpeechProcessor", LJSpeechProcessor),
+        ("KSSProcessor", KSSProcessor),
+        ("BakerProcessor", BakerProcessor),
+        ("LibriTTSProcessor", LibriTTSProcessor),
+    ]
+)
+
+
+class AutoProcessor:
+    def __init__(self):
+        raise EnvironmentError(
+            "AutoProcessor is designed to be instantiated "
+            "using the `AutoProcessor.from_pretrained(pretrained_path)` method."
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_path, **kwargs):
+        with open(pretrained_path, "r") as f:
+            config = json.load(f)
+
+        try:
+            processor_name = config["processor_name"]
+            processor_class = CONFIG_MAPPING[processor_name]
+            processor_class = processor_class(
+                data_dir=None, loaded_mapper_path=pretrained_path
+            )
+            return processor_class
+        except Exception:
+            raise ValueError(
+                "Unrecognized processor in {}. "
+                "Should have a `processor_name` key in its config.json, or contain one of the following strings "
+                "in its name: {}".format(
+                    pretrained_path, ", ".join(CONFIG_MAPPING.keys())
+                )
+            )

test/files/mapper.json

@@ -13,5 +13,5 @@
     "1": "b",
     "2": "@ph"
   },
-  "processor_name": "TestProcessor"
+  "processor_name": "KSSProcessor"
 }

test/test_auto_model.py test/test_auto.pytest/test_auto_model.py renamed to test/test_auto.py

@@ -20,6 +20,7 @@
 import tensorflow as tf
 
 from tensorflow_tts.inference import AutoConfig
+from tensorflow_tts.inference import AutoProcessor
 from tensorflow_tts.inference import TFAutoModel
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -49,3 +50,9 @@
 def test_auto_model(config_path):
     config = AutoConfig.from_pretrained(pretrained_path=config_path)
     model = TFAutoModel.from_pretrained(config=config, pretrained_path=None)
+
+
+@pytest.fixture
+def test_auto_processor():
+    processor = AutoProcessor(data_dir=None, loaded_mapper_path="./files/mapper.json")
+    return processor