This commit is contained in:
litagin02
2024-03-15 12:27:46 +09:00
parent 867855ace9
commit 3d8b60c03c
7 changed files with 243 additions and 199 deletions

View File

@@ -1,93 +1,61 @@
import argparse
import json import json
import os
from collections import defaultdict from collections import defaultdict
from pathlib import Path
from random import shuffle from random import shuffle
from typing import Optional from typing import Optional
import click
from tqdm import tqdm from tqdm import tqdm
from config import config from config import Preprocess_text_config, config
from style_bert_vits2.logging import logger from style_bert_vits2.logging import logger
from style_bert_vits2.nlp import clean_text from style_bert_vits2.nlp import clean_text
from style_bert_vits2.nlp.japanese import pyopenjtalk_worker from style_bert_vits2.nlp.japanese import pyopenjtalk_worker
from style_bert_vits2.nlp.japanese.user_dict import update_dict from style_bert_vits2.nlp.japanese.user_dict import update_dict
from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
# このプロセスからはワーカーを起動して辞書を使いたいので、ここで初期化 # このプロセスからはワーカーを起動して辞書を使いたいので、ここで初期化
pyopenjtalk_worker.initialize_worker() pyopenjtalk_worker.initialize_worker()
# dict_data/ 以下の辞書データを pyopenjtalk に適用 # dict_data/ 以下の辞書データを pyopenjtalk に適用
update_dict() update_dict()
preprocess_text_config = config.preprocess_text_config
preprocess_text_config: Preprocess_text_config = config.preprocess_text_config
# Count lines for tqdm # Count lines for tqdm
def count_lines(file_path: str): def count_lines(file_path: Path):
with open(file_path, "r", encoding="utf-8") as file: with file_path.open("r", encoding="utf-8") as file:
return sum(1 for _ in file) return sum(1 for _ in file)
@click.command() def write_error_log(error_log_path: Path, line: str, error: Exception):
@click.option( with error_log_path.open("a", encoding="utf-8") as error_log:
"--transcription-path", error_log.write(f"{line.strip()}\n{error}\n\n")
default=preprocess_text_config.transcription_path,
type=click.Path(exists=True, file_okay=True, dir_okay=False),
) def process_line(
@click.option("--cleaned-path", default=preprocess_text_config.cleaned_path) line: str,
@click.option("--train-path", default=preprocess_text_config.train_path) transcription_path: Path,
@click.option("--val-path", default=preprocess_text_config.val_path) correct_path: bool,
@click.option(
"--config-path",
default=preprocess_text_config.config_path,
type=click.Path(exists=True, file_okay=True, dir_okay=False),
)
@click.option("--val-per-lang", default=preprocess_text_config.val_per_lang)
@click.option("--max-val-total", default=preprocess_text_config.max_val_total)
@click.option("--clean/--no-clean", default=preprocess_text_config.clean)
@click.option("-y", "--yml_config")
@click.option("--use_jp_extra", is_flag=True)
@click.option("--yomi_error", default="raise")
def preprocess(
transcription_path: str,
cleaned_path: Optional[str],
train_path: str,
val_path: str,
config_path: str,
val_per_lang: int,
max_val_total: int,
clean: bool,
yml_config: str, # 这个不要删
use_jp_extra: bool, use_jp_extra: bool,
yomi_error: str, yomi_error: str,
): ):
assert yomi_error in ["raise", "skip", "use"] splitted_line = line.strip().split("|")
if cleaned_path == "" or cleaned_path is None: if len(splitted_line) != 4:
cleaned_path = transcription_path + ".cleaned" raise ValueError(f"Invalid line format: {line.strip()}")
utt, spk, language, text = splitted_line
error_log_path = os.path.join(os.path.dirname(cleaned_path), "text_error.log")
if os.path.exists(error_log_path):
os.remove(error_log_path)
error_count = 0
if clean:
total_lines = count_lines(transcription_path)
with open(cleaned_path, "w", encoding="utf-8") as out_file:
with open(transcription_path, "r", encoding="utf-8") as trans_file:
for line in tqdm(trans_file, file=SAFE_STDOUT, total=total_lines):
try:
utt, spk, language, text = line.strip().split("|")
norm_text, phones, tones, word2ph = clean_text( norm_text, phones, tones, word2ph = clean_text(
text=text, text=text,
language=language, # type: ignore language=language, # type: ignore
use_jp_extra=use_jp_extra, use_jp_extra=use_jp_extra,
raise_yomi_error=(yomi_error != "use"), raise_yomi_error=(yomi_error != "use"),
) )
if correct_path:
utt = str(transcription_path.parent / "wavs" / utt)
out_file.write( return "{}|{}|{}|{}|{}|{}|{}\n".format(
"{}|{}|{}|{}|{}|{}|{}\n".format(
utt, utt,
spk, spk,
language, language,
@@ -96,50 +64,98 @@ def preprocess(
" ".join([str(i) for i in tones]), " ".join([str(i) for i in tones]),
" ".join([str(i) for i in word2ph]), " ".join([str(i) for i in word2ph]),
) )
def preprocess(
transcription_path: Path,
cleaned_path: Optional[Path],
train_path: Path,
val_path: Path,
config_path: Path,
val_per_lang: int,
max_val_total: int,
# clean: bool,
use_jp_extra: bool,
yomi_error: str,
correct_path: bool,
):
assert yomi_error in ["raise", "skip", "use"]
if cleaned_path == "" or cleaned_path is None:
cleaned_path = transcription_path.with_name(
transcription_path.name + ".cleaned"
) )
error_log_path = transcription_path.parent / "text_error.log"
if error_log_path.exists():
error_log_path.unlink()
error_count = 0
total_lines = count_lines(transcription_path)
# transcription_path から 1行ずつ読み込んで文章処理して cleaned_path に書き込む
with (
transcription_path.open("r", encoding="utf-8") as trans_file,
cleaned_path.open("w", encoding="utf-8") as out_file,
):
for line in tqdm(trans_file, file=SAFE_STDOUT, total=total_lines):
try:
processed_line = process_line(
line,
transcription_path,
correct_path,
use_jp_extra,
yomi_error,
)
out_file.write(processed_line)
except Exception as e: except Exception as e:
logger.error( logger.error(
f"An error occurred at line:\n{line.strip()}\n{e}", f"An error occurred at line:\n{line.strip()}\n{e}", encoding="utf-8"
encoding="utf-8",
) )
with open(error_log_path, "a", encoding="utf-8") as error_log: write_error_log(error_log_path, line, e)
error_log.write(f"{line.strip()}\n{e}\n\n")
error_count += 1 error_count += 1
transcription_path = cleaned_path transcription_path = cleaned_path
spk_utt_map = defaultdict(list)
spk_id_map = {}
current_sid = 0
with open(transcription_path, "r", encoding="utf-8") as f: # 各話者ごとのlineの辞書
audioPaths = set() spk_utt_map: dict[str, list[str]] = defaultdict(list)
countSame = 0
countNotFound = 0 # 話者からIDへの写像
spk_id_map: dict[str, int] = {}
# 話者ID
current_sid: int = 0
# 音源ファイルのチェックや、spk_id_mapの作成
with transcription_path.open("r", encoding="utf-8") as f:
audio_paths: set[str] = set()
count_same = 0
count_not_found = 0
for line in f.readlines(): for line in f.readlines():
utt, spk, language, text, phones, tones, word2ph = line.strip().split("|") utt, spk = line.strip().split("|")[:2]
if utt in audioPaths: if utt in audio_paths:
# 过滤数据集错误相同的音频匹配多个文本导致后续bert出问题 logger.warning(f"Same audio file appears multiple times: {utt}")
logger.warning(f"Same audio matches multiple texts: {line}") count_same += 1
countSame += 1
continue continue
if not os.path.isfile(utt): if not Path(utt).is_file():
# 过滤数据集错误:不存在对应音频
logger.warning(f"Audio not found: {utt}") logger.warning(f"Audio not found: {utt}")
countNotFound += 1 count_not_found += 1
continue continue
audioPaths.add(utt) audio_paths.add(utt)
spk_utt_map[language].append(line) spk_utt_map[spk].append(line)
# 新しい話者が出てきたら話者IDを割り当て、current_sidを1増やす
if spk not in spk_id_map.keys(): if spk not in spk_id_map.keys():
spk_id_map[spk] = current_sid spk_id_map[spk] = current_sid
current_sid += 1 current_sid += 1
if countSame > 0 or countNotFound > 0: if count_same > 0 or count_not_found > 0:
logger.warning( logger.warning(
f"Total repeated audios: {countSame}, Total number of audio not found: {countNotFound}" f"Total repeated audios: {count_same}, Total number of audio not found: {count_not_found}"
) )
train_list = [] train_list: list[str] = []
val_list = [] val_list: list[str] = []
# 各話者ごとにシャッフルして、val_per_lang個をval_listに、残りをtrain_listに追加
for spk, utts in spk_utt_map.items(): for spk, utts in spk_utt_map.items():
shuffle(utts) shuffle(utts)
val_list += utts[:val_per_lang] val_list += utts[:val_per_lang]
@@ -150,26 +166,21 @@ def preprocess(
train_list += val_list[max_val_total:] train_list += val_list[max_val_total:]
val_list = val_list[:max_val_total] val_list = val_list[:max_val_total]
with open(train_path, "w", encoding="utf-8") as f: with train_path.open("w", encoding="utf-8") as f:
for line in train_list: for line in train_list:
f.write(line) f.write(line)
with open(val_path, "w", encoding="utf-8") as f: with val_path.open("w", encoding="utf-8") as f:
for line in val_list: for line in val_list:
f.write(line) f.write(line)
json_config = json.load(open(config_path, encoding="utf-8")) with config_path.open("r", encoding="utf-8") as f:
json_config = json.load(f)
json_config["data"]["spk2id"] = spk_id_map json_config["data"]["spk2id"] = spk_id_map
json_config["data"]["n_speakers"] = len(spk_id_map) json_config["data"]["n_speakers"] = len(spk_id_map)
# 新增写入:写入训练版本、数据集路径
# json_config["version"] = latest_version with config_path.open("w", encoding="utf-8") as f:
json_config["data"]["training_files"] = os.path.normpath(train_path).replace(
"\\", "/"
)
json_config["data"]["validation_files"] = os.path.normpath(val_path).replace(
"\\", "/"
)
with open(config_path, "w", encoding="utf-8") as f:
json.dump(json_config, f, indent=2, ensure_ascii=False) json.dump(json_config, f, indent=2, ensure_ascii=False)
if error_count > 0: if error_count > 0:
if yomi_error == "skip": if yomi_error == "skip":
@@ -194,4 +205,50 @@ def preprocess(
if __name__ == "__main__": if __name__ == "__main__":
preprocess() parser = argparse.ArgumentParser()
parser.add_argument(
"--transcription-path", default=preprocess_text_config.transcription_path
)
parser.add_argument("--cleaned-path", default=preprocess_text_config.cleaned_path)
parser.add_argument("--train-path", default=preprocess_text_config.train_path)
parser.add_argument("--val-path", default=preprocess_text_config.val_path)
parser.add_argument("--config-path", default=preprocess_text_config.config_path)
# 「話者ごと」のバリデーションデータ数、言語ごとではない!
# 元のコードや設定ファイルでval_per_langとなっていたので名前をそのままにしている
parser.add_argument(
"--val-per-lang",
default=preprocess_text_config.val_per_lang,
help="Number of validation data per SPEAKER, not per language (due to compatibility with the original code).",
)
parser.add_argument("--max-val-total", default=preprocess_text_config.max_val_total)
parser.add_argument("--use_jp_extra", action="store_true")
parser.add_argument("--yomi_error", default="raise")
parser.add_argument("--correct_path", action="store_true")
args = parser.parse_args()
logger.debug(f"args: {args}")
transcription_path = Path(args.transcription_path)
cleaned_path = Path(args.cleaned_path) if args.cleaned_path else None
train_path = Path(args.train_path)
val_path = Path(args.val_path)
config_path = Path(args.config_path)
val_per_lang = int(args.val_per_lang)
max_val_total = int(args.max_val_total)
use_jp_extra: bool = args.use_jp_extra
yomi_error: str = args.yomi_error
correct_path: bool = args.correct_path
preprocess(
transcription_path=transcription_path,
cleaned_path=cleaned_path,
train_path=train_path,
val_path=val_path,
config_path=config_path,
val_per_lang=val_per_lang,
max_val_total=max_val_total,
use_jp_extra=use_jp_extra,
yomi_error=yomi_error,
correct_path=correct_path,
)

View File

@@ -14,6 +14,7 @@ from config import config
from style_bert_vits2.logging import logger from style_bert_vits2.logging import logger
from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
DEFAULT_BLOCK_SIZE: float = 0.400 # seconds DEFAULT_BLOCK_SIZE: float = 0.400 # seconds
@@ -33,6 +34,9 @@ def normalize_audio(data: NDArray[Any], sr: int):
def resample(file: Path, output_dir: Path, target_sr: int, normalize: bool, trim: bool): def resample(file: Path, output_dir: Path, target_sr: int, normalize: bool, trim: bool):
"""
fileを読み込んで、target_srなwavファイルに変換してoutput_dir直下に保存する
"""
try: try:
# librosaが読めるファイルかチェック # librosaが読めるファイルかチェック
# wav以外にもmp3やoggやflacなども読める # wav以外にもmp3やoggやflacなども読める

View File

@@ -12,5 +12,4 @@ logger.add(
format="<g>{time:MM-DD HH:mm:ss}</g> |<lvl>{level:^8}</lvl>| {file}:{line} | {message}", format="<g>{time:MM-DD HH:mm:ss}</g> |<lvl>{level:^8}</lvl>| {file}:{line} | {message}",
backtrace=True, backtrace=True,
diagnose=True, diagnose=True,
level="TRACE",
) )

View File

@@ -36,7 +36,7 @@ def normalize_text(text: str) -> str:
res = unicodedata.normalize("NFKC", text) # ここでアルファベットは半角になる res = unicodedata.normalize("NFKC", text) # ここでアルファベットは半角になる
res = __convert_numbers_to_words(res) # 「100円」→「百円」等 res = __convert_numbers_to_words(res) # 「100円」→「百円」等
# 「~」と「〜」と「~」も長音記号として扱う # 「~」と「〜」と「~」も長音記号として扱う
res = res.replace("~", "") res = res.replace("~", "")
res = res.replace("", "") res = res.replace("", "")
res = res.replace("", "") res = res.replace("", "")

View File

@@ -1,5 +1,4 @@
import argparse import argparse
import warnings
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from typing import Any from typing import Any

View File

@@ -6,12 +6,12 @@ import gradio as gr
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import numpy as np import numpy as np
import yaml import yaml
from config import config
from scipy.spatial.distance import pdist, squareform from scipy.spatial.distance import pdist, squareform
from sklearn.cluster import DBSCAN, AgglomerativeClustering, KMeans from sklearn.cluster import DBSCAN, AgglomerativeClustering, KMeans
from sklearn.manifold import TSNE from sklearn.manifold import TSNE
from umap import UMAP from umap import UMAP
from config import config
from style_bert_vits2.constants import DEFAULT_STYLE, GRADIO_THEME from style_bert_vits2.constants import DEFAULT_STYLE, GRADIO_THEME
from style_bert_vits2.logging import logger from style_bert_vits2.logging import logger

View File

@@ -24,32 +24,31 @@ tensorboard_executed = False
# Get path settings # Get path settings
with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f: with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
path_config: dict[str, str] = yaml.safe_load(f.read()) path_config: dict[str, str] = yaml.safe_load(f.read())
dataset_root = path_config["dataset_root"] dataset_root = Path(path_config["dataset_root"])
# assets_root = path_config["assets_root"]
def get_path(model_name): def get_path(model_name: str) -> tuple[Path, Path, Path, Path, Path]:
assert model_name != "", "モデル名は空にできません" assert model_name != "", "モデル名は空にできません"
dataset_path = os.path.join(dataset_root, model_name) dataset_path = dataset_root / model_name
lbl_path = os.path.join(dataset_path, "esd.list") lbl_path = dataset_path / "esd.list"
train_path = os.path.join(dataset_path, "train.list") train_path = dataset_path / "train.list"
val_path = os.path.join(dataset_path, "val.list") val_path = dataset_path / "val.list"
config_path = os.path.join(dataset_path, "config.json") config_path = dataset_path / "config.json"
return dataset_path, lbl_path, train_path, val_path, config_path return dataset_path, lbl_path, train_path, val_path, config_path
def initialize( def initialize(
model_name, model_name: str,
batch_size, batch_size: int,
epochs, epochs: int,
save_every_steps, save_every_steps: int,
freeze_EN_bert, freeze_EN_bert: bool,
freeze_JP_bert, freeze_JP_bert: bool,
freeze_ZH_bert, freeze_ZH_bert: bool,
freeze_style, freeze_style: bool,
freeze_decoder, freeze_decoder: bool,
use_jp_extra, use_jp_extra: bool,
log_interval, log_interval: int,
): ):
global logger_handler global logger_handler
dataset_path, _, train_path, val_path, config_path = get_path(model_name) dataset_path, _, train_path, val_path, config_path = get_path(model_name)
@@ -72,8 +71,8 @@ def initialize(
with open(default_config_path, "r", encoding="utf-8") as f: with open(default_config_path, "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
config["model_name"] = model_name config["model_name"] = model_name
config["data"]["training_files"] = train_path config["data"]["training_files"] = str(train_path)
config["data"]["validation_files"] = val_path config["data"]["validation_files"] = str(val_path)
config["train"]["batch_size"] = batch_size config["train"]["batch_size"] = batch_size
config["train"]["epochs"] = epochs config["train"]["epochs"] = epochs
config["train"]["eval_interval"] = save_every_steps config["train"]["eval_interval"] = save_every_steps
@@ -90,18 +89,18 @@ def initialize(
# 今はデフォルトであるが、以前は非JP-Extra版になくバグの原因になるので念のため # 今はデフォルトであるが、以前は非JP-Extra版になくバグの原因になるので念のため
config["data"]["use_jp_extra"] = use_jp_extra config["data"]["use_jp_extra"] = use_jp_extra
model_path = os.path.join(dataset_path, "models") model_path = dataset_path / "models"
if os.path.exists(model_path): if model_path.exists():
logger.warning( logger.warning(
f"Step 1: {model_path} already exists, so copy it to backup to {model_path}_backup" f"Step 1: {model_path} already exists, so copy it to backup to {model_path}_backup"
) )
shutil.copytree( shutil.copytree(
src=model_path, src=model_path,
dst=os.path.join(dataset_path, "models_backup"), dst=dataset_path / "models_backup",
dirs_exist_ok=True, dirs_exist_ok=True,
) )
shutil.rmtree(model_path) shutil.rmtree(model_path)
pretrained_dir = "pretrained" if not use_jp_extra else "pretrained_jp_extra" pretrained_dir = Path("pretrained" if not use_jp_extra else "pretrained_jp_extra")
try: try:
shutil.copytree( shutil.copytree(
src=pretrained_dir, src=pretrained_dir,
@@ -113,30 +112,29 @@ def initialize(
with open(config_path, "w", encoding="utf-8") as f: with open(config_path, "w", encoding="utf-8") as f:
json.dump(config, f, indent=2, ensure_ascii=False) json.dump(config, f, indent=2, ensure_ascii=False)
if not os.path.exists("config.yml"): if not Path("config.yml").exists():
shutil.copy(src="default_config.yml", dst="config.yml") shutil.copy(src="default_config.yml", dst="config.yml")
# yml_data = safe_load(open("config.yml", "r", encoding="utf-8"))
with open("config.yml", "r", encoding="utf-8") as f: with open("config.yml", "r", encoding="utf-8") as f:
yml_data = yaml.safe_load(f) yml_data = yaml.safe_load(f)
yml_data["model_name"] = model_name yml_data["model_name"] = model_name
yml_data["dataset_path"] = dataset_path yml_data["dataset_path"] = str(dataset_path)
with open("config.yml", "w", encoding="utf-8") as f: with open("config.yml", "w", encoding="utf-8") as f:
yaml.dump(yml_data, f, allow_unicode=True) yaml.dump(yml_data, f, allow_unicode=True)
logger.success("Step 1: initialization finished.") logger.success("Step 1: initialization finished.")
return True, "Step 1, Success: 初期設定が完了しました" return True, "Step 1, Success: 初期設定が完了しました"
def resample(model_name, normalize, trim, num_processes): def resample(model_name: str, normalize: bool, trim: bool, num_processes: int):
logger.info("Step 2: start resampling...") logger.info("Step 2: start resampling...")
dataset_path, _, _, _, _ = get_path(model_name) dataset_path, _, _, _, _ = get_path(model_name)
input_dir = os.path.join(dataset_path, "raw") input_dir = dataset_path / "raw"
output_dir = os.path.join(dataset_path, "wavs") output_dir = dataset_path / "wavs"
cmd = [ cmd = [
"resample.py", "resample.py",
"-i", "-i",
input_dir, str(input_dir),
"-o", "-o",
output_dir, str(output_dir),
"--num_processes", "--num_processes",
str(num_processes), str(num_processes),
"--sr", "--sr",
@@ -157,42 +155,30 @@ def resample(model_name, normalize, trim, num_processes):
return True, "Step 2, Success: 音声ファイルの前処理が完了しました" return True, "Step 2, Success: 音声ファイルの前処理が完了しました"
def preprocess_text(model_name, use_jp_extra, val_per_lang, yomi_error): def preprocess_text(
model_name: str, use_jp_extra: bool, val_per_lang: int, yomi_error: str
):
logger.info("Step 3: start preprocessing text...") logger.info("Step 3: start preprocessing text...")
dataset_path, lbl_path, train_path, val_path, config_path = get_path(model_name) _, lbl_path, train_path, val_path, config_path = get_path(model_name)
try: if not lbl_path.exists():
lines = open(lbl_path, "r", encoding="utf-8").readlines()
except FileNotFoundError:
logger.error(f"Step 3: {lbl_path} not found.") logger.error(f"Step 3: {lbl_path} not found.")
return False, f"Step 3, Error: 書き起こしファイル {lbl_path} が見つかりません。" return False, f"Step 3, Error: 書き起こしファイル {lbl_path} が見つかりません。"
new_lines = []
for line in lines:
if len(line.strip().split("|")) != 4:
logger.error(f"Step 3: {lbl_path} has invalid format at line:\n{line}")
return (
False,
f"Step 3, Error: 書き起こしファイル次の行の形式が不正です:\n{line}",
)
path, spk, language, text = line.strip().split("|")
# pathをファイル名だけ取り出して正しいパスに変更
path = Path(dataset_path) / "wavs" / Path(path).name
new_lines.append(f"{path}|{spk}|{language}|{text}\n")
with open(lbl_path, "w", encoding="utf-8") as f:
f.writelines(new_lines)
cmd = [ cmd = [
"preprocess_text.py", "preprocess_text.py",
"--config-path", "--config-path",
config_path, str(config_path),
"--transcription-path", "--transcription-path",
lbl_path, str(lbl_path),
"--train-path", "--train-path",
train_path, str(train_path),
"--val-path", "--val-path",
val_path, str(val_path),
"--val-per-lang", "--val-per-lang",
str(val_per_lang), str(val_per_lang),
"--yomi_error", "--yomi_error",
yomi_error, yomi_error,
"--correct_path", # 音声ファイルのパスを正しいパスに修正する
] ]
if use_jp_extra: if use_jp_extra:
cmd.append("--use_jp_extra") cmd.append("--use_jp_extra")
@@ -213,17 +199,11 @@ def preprocess_text(model_name, use_jp_extra, val_per_lang, yomi_error):
return True, "Step 3, Success: 書き起こしファイルの前処理が完了しました" return True, "Step 3, Success: 書き起こしファイルの前処理が完了しました"
def bert_gen(model_name): def bert_gen(model_name: str):
logger.info("Step 4: start bert_gen...") logger.info("Step 4: start bert_gen...")
_, _, _, _, config_path = get_path(model_name) _, _, _, _, config_path = get_path(model_name)
success, message = run_script_with_log( success, message = run_script_with_log(
[ ["bert_gen.py", "--config", str(config_path)]
"bert_gen.py",
"--config",
config_path,
# "--num_processes", # bert_genは重いのでプロセス数いじらない
# str(num_processes),
]
) )
if not success: if not success:
logger.error("Step 4: bert_gen failed.") logger.error("Step 4: bert_gen failed.")
@@ -238,14 +218,14 @@ def bert_gen(model_name):
return True, "Step 4, Success: BERT特徴ファイルの生成が完了しました" return True, "Step 4, Success: BERT特徴ファイルの生成が完了しました"
def style_gen(model_name, num_processes): def style_gen(model_name: str, num_processes: int):
logger.info("Step 5: start style_gen...") logger.info("Step 5: start style_gen...")
_, _, _, _, config_path = get_path(model_name) _, _, _, _, config_path = get_path(model_name)
success, message = run_script_with_log( success, message = run_script_with_log(
[ [
"style_gen.py", "style_gen.py",
"--config", "--config",
config_path, str(config_path),
"--num_processes", "--num_processes",
str(num_processes), str(num_processes),
] ]
@@ -267,22 +247,22 @@ def style_gen(model_name, num_processes):
def preprocess_all( def preprocess_all(
model_name, model_name: str,
batch_size, batch_size: int,
epochs, epochs: int,
save_every_steps, save_every_steps: int,
num_processes, num_processes: int,
normalize, normalize: bool,
trim, trim: bool,
freeze_EN_bert, freeze_EN_bert: bool,
freeze_JP_bert, freeze_JP_bert: bool,
freeze_ZH_bert, freeze_ZH_bert: bool,
freeze_style, freeze_style: bool,
freeze_decoder, freeze_decoder: bool,
use_jp_extra, use_jp_extra: bool,
val_per_lang, val_per_lang: int,
log_interval, log_interval: int,
yomi_error, yomi_error: str,
): ):
if model_name == "": if model_name == "":
return False, "Error: モデル名を入力してください" return False, "Error: モデル名を入力してください"
@@ -333,18 +313,23 @@ def preprocess_all(
) )
def train(model_name, skip_style=False, use_jp_extra=True, speedup=False): def train(
model_name: str,
skip_style: bool = False,
use_jp_extra: bool = True,
speedup: bool = False,
):
dataset_path, _, _, _, config_path = get_path(model_name) dataset_path, _, _, _, config_path = get_path(model_name)
# 学習再開の場合念のためconfig.ymlの名前等を更新 # 学習再開の場合を考えて念のためconfig.ymlの名前等を更新
with open("config.yml", "r", encoding="utf-8") as f: with open("config.yml", "r", encoding="utf-8") as f:
yml_data = yaml.safe_load(f) yml_data = yaml.safe_load(f)
yml_data["model_name"] = model_name yml_data["model_name"] = model_name
yml_data["dataset_path"] = dataset_path yml_data["dataset_path"] = str(dataset_path)
with open("config.yml", "w", encoding="utf-8") as f: with open("config.yml", "w", encoding="utf-8") as f:
yaml.dump(yml_data, f, allow_unicode=True) yaml.dump(yml_data, f, allow_unicode=True)
train_py = "train_ms.py" if not use_jp_extra else "train_ms_jp_extra.py" train_py = "train_ms.py" if not use_jp_extra else "train_ms_jp_extra.py"
cmd = [train_py, "--config", config_path, "--model", dataset_path] cmd = [train_py, "--config", str(config_path), "--model", str(dataset_path)]
if skip_style: if skip_style:
cmd.append("--skip_default_style") cmd.append("--skip_default_style")
if speedup: if speedup:
@@ -360,7 +345,7 @@ def train(model_name, skip_style=False, use_jp_extra=True, speedup=False):
return True, "Success: 学習が完了しました" return True, "Success: 学習が完了しました"
def wait_for_tensorboard(port=6006, timeout=10): def wait_for_tensorboard(port: int = 6006, timeout: float = 10):
start_time = time.time() start_time = time.time()
while True: while True:
try: try:
@@ -375,7 +360,7 @@ def wait_for_tensorboard(port=6006, timeout=10):
time.sleep(0.1) time.sleep(0.1)
def run_tensorboard(model_name): def run_tensorboard(model_name: str):
global tensorboard_executed global tensorboard_executed
if not tensorboard_executed: if not tensorboard_executed:
python = sys.executable python = sys.executable