Change args from input output to model_name

This commit is contained in:
litagin02
2024-02-19 19:36:14 +09:00
parent 2053ddd101
commit b7ffc37ba9
2 changed files with 16 additions and 22 deletions

View File

@@ -2,6 +2,7 @@ import argparse
import os import os
import sys import sys
import yaml
from faster_whisper import WhisperModel from faster_whisper import WhisperModel
from tqdm import tqdm from tqdm import tqdm
@@ -20,25 +21,29 @@ def transcribe(wav_path, initial_prompt=None, language="ja"):
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--input_dir", "-i", type=str, default="raw") parser.add_argument("--model_name", type=str, required=True)
parser.add_argument("--output_file", "-o", type=str, default="esd.list")
parser.add_argument( parser.add_argument(
"--initial_prompt", type=str, default="こんにちは。元気、ですかー?ふふっ、私は……ちゃんと元気だよ!" "--initial_prompt",
type=str,
default="こんにちは。元気、ですかー?ふふっ、私は……ちゃんと元気だよ!",
) )
parser.add_argument( parser.add_argument(
"--language", type=str, default="ja", choices=["ja", "en", "zh"] "--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("--model", type=str, default="large-v3")
parser.add_argument("--device", type=str, default="cuda") parser.add_argument("--device", type=str, default="cuda")
parser.add_argument("--compute_type", type=str, default="bfloat16") parser.add_argument("--compute_type", type=str, default="bfloat16")
args = parser.parse_args() 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 model_name = args.model_name
output_file = args.output_file
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 initial_prompt = args.initial_prompt
language = args.language language = args.language
device = args.device device = args.device
@@ -79,6 +84,6 @@ if __name__ == "__main__":
text = transcribe( text = transcribe(
wav_file, initial_prompt=initial_prompt, language=language 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() f.flush()
sys.exit(0) sys.exit(0)

View File

@@ -46,24 +46,17 @@ def do_slice(
def do_transcribe( 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 == "": if model_name == "":
return "Error: モデル名を入力してください。" return "Error: モデル名を入力してください。"
if initial_prompt == "": if initial_prompt == "":
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( success, message = run_script_with_log(
[ [
"transcribe.py", "transcribe.py",
"--input_dir", "--model_name",
input_dir,
"--output_file",
output_file,
"--speaker_name",
model_name, model_name,
"--model", "--model",
whisper_model, whisper_model,
@@ -153,9 +146,6 @@ with gr.Blocks(theme=GRADIO_THEME) as app:
result1 = gr.Textbox(label="結果") result1 = gr.Textbox(label="結果")
with gr.Row(): with gr.Row():
with gr.Column(): with gr.Column():
raw_dir = gr.Textbox(
label="書き起こしたい音声ファイルが入っているフォルダ(スライスした場合など、`Data/{モデル名}/raw`の場合は省略可",
)
whisper_model = gr.Dropdown( whisper_model = gr.Dropdown(
["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"], ["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"],
label="Whisperモデル", label="Whisperモデル",
@@ -197,7 +187,6 @@ with gr.Blocks(theme=GRADIO_THEME) as app:
compute_type, compute_type,
language, language,
initial_prompt, initial_prompt,
raw_dir,
device, device,
], ],
outputs=[result2], outputs=[result2],