From d261a63afbb292ee0a6a97025a7107af556e2604 Mon Sep 17 00:00:00 2001
From: Michael Hansen <mike@rhasspy.org>
Date: Tue, 12 Sep 2023 14:11:08 -0500
Subject: [PATCH] Add extra config items

---
 src/python/piper_train/preprocess.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/src/python/piper_train/preprocess.py b/src/python/piper_train/preprocess.py
index b89da57..9c23d89 100644
--- a/src/python/piper_train/preprocess.py
+++ b/src/python/piper_train/preprocess.py
@@ -82,6 +82,21 @@ def main() -> None:
         help="Casing applied to utterance text",
     )
     #
+    parser.add_argument(
+        "--dataset-name",
+        help="Name of dataset to put in config (default: name of <ouput_dir>/../)",
+    )
+    parser.add_argument(
+        "--audio-quality",
+        help="Audio quality to put in config (default: name of <output_dir>)",
+    )
+    #
+    parser.add_argument(
+        "--tashkeel",
+        action="store_true",
+        help="Diacritize Arabic text with libtashkeel",
+    )
+    #
     parser.add_argument(
         "--skip-audio", action="store_true", help="Don't preprocess audio"
     )
@@ -147,15 +162,23 @@ def main() -> None:
         _LOGGER.info("Single speaker dataset")
 
     # Write config
+    audio_quality = args.audio_quality or args.output_dir.name
+    dataset_name = args.dataset_name or args.output_dir.parent.name
+
     with open(args.output_dir / "config.json", "w", encoding="utf-8") as config_file:
         json.dump(
             {
+                "dataset": dataset_name,
                 "audio": {
                     "sample_rate": args.sample_rate,
+                    "quality": audio_quality,
                 },
                 "espeak": {
                     "voice": args.language,
                 },
+                "language": {
+                    "code": args.language,
+                },
                 "inference": {"noise_scale": 0.667, "length_scale": 1, "noise_w": 0.8},
                 "phoneme_type": args.phoneme_type.value,
                 "phoneme_map": {},