From b7ffc37ba98d10f5415fa4d027c5fc9fedb53ea8 Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Mon, 19 Feb 2024 19:36:14 +0900
Subject: [PATCH] Change args from input output to model_name

---
 transcribe.py    | 21 +++++++++++++--------
 webui_dataset.py | 17 +++--------------
 2 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/transcribe.py b/transcribe.py
index d9fb98b..0ec9e99 100644
--- a/transcribe.py
+++ b/transcribe.py
@@ -2,6 +2,7 @@ import argparse
 import os
 import sys
 
+import yaml
 from faster_whisper import WhisperModel
 from tqdm import tqdm
 
@@ -20,25 +21,29 @@ def transcribe(wav_path, initial_prompt=None, language="ja"):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--input_dir", "-i", type=str, default="raw")
-    parser.add_argument("--output_file", "-o", type=str, default="esd.list")
+    parser.add_argument("--model_name", type=str, required=True)
     parser.add_argument(
-        "--initial_prompt", type=str, default="こんにちは。元気、ですかー？ふふっ、私は……ちゃんと元気だよ！"
+        "--initial_prompt",
+        type=str,
+        default="こんにちは。元気、ですかー？ふふっ、私は……ちゃんと元気だよ！",
     )
     parser.add_argument(
         "--language", type=str, default="ja", choices=["ja", "en", "zh"]
     )
-    parser.add_argument("--speaker_name", type=str, required=True)
     parser.add_argument("--model", type=str, default="large-v3")
     parser.add_argument("--device", type=str, default="cuda")
     parser.add_argument("--compute_type", type=str, default="bfloat16")
 
     args = parser.parse_args()
 
-    speaker_name = args.speaker_name
+    with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
+        path_config: dict[str, str] = yaml.safe_load(f.read())
+        dataset_root = path_config["dataset_root"]
 
-    input_dir = args.input_dir
-    output_file = args.output_file
+    model_name = args.model_name
+
+    input_dir = os.path.join(dataset_root, model_name, "raw")
+    output_file = os.path.join(dataset_root, model_name, "esd.list")
     initial_prompt = args.initial_prompt
     language = args.language
     device = args.device
@@ -79,6 +84,6 @@ if __name__ == "__main__":
             text = transcribe(
                 wav_file, initial_prompt=initial_prompt, language=language
             )
-            f.write(f"{file_name}|{speaker_name}|{language_id}|{text}\n")
+            f.write(f"{file_name}|{model_name}|{language_id}|{text}\n")
             f.flush()
     sys.exit(0)
diff --git a/webui_dataset.py b/webui_dataset.py
index 02fb15d..a795575 100644
--- a/webui_dataset.py
+++ b/webui_dataset.py
@@ -46,24 +46,17 @@ def do_slice(
 
 
 def do_transcribe(
-    model_name, whisper_model, compute_type, language, initial_prompt, input_dir, device
+    model_name, whisper_model, compute_type, language, initial_prompt, device
 ):
     if model_name == "":
         return "Error: モデル名を入力してください。"
     if initial_prompt == "":
         initial_prompt = "こんにちは。元気、ですかー？私は……ふふっ、ちゃんと元気だよ！"
-    # logger.debug(f"initial_prompt: {initial_prompt}")
-    if input_dir == "":
-        input_dir = os.path.join(dataset_root, model_name, "raw")
-    output_file = os.path.join(dataset_root, model_name, "esd.list")
+
     success, message = run_script_with_log(
         [
             "transcribe.py",
-            "--input_dir",
-            input_dir,
-            "--output_file",
-            output_file,
-            "--speaker_name",
+            "--model_name",
             model_name,
             "--model",
             whisper_model,
@@ -153,9 +146,6 @@ with gr.Blocks(theme=GRADIO_THEME) as app:
             result1 = gr.Textbox(label="結果")
     with gr.Row():
         with gr.Column():
-            raw_dir = gr.Textbox(
-                label="書き起こしたい音声ファイルが入っているフォルダ（スライスした場合など、`Data/{モデル名}/raw`の場合は省略可",
-            )
             whisper_model = gr.Dropdown(
                 ["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"],
                 label="Whisperモデル",
@@ -197,7 +187,6 @@ with gr.Blocks(theme=GRADIO_THEME) as app:
             compute_type,
             language,
             initial_prompt,
-            raw_dir,
             device,
         ],
         outputs=[result2],