fix

gradio_tabs/train.py

@@ -58,6 +58,7 @@ def initialize(
     freeze_decoder: bool,
     use_jp_extra: bool,
     log_interval: int,
+    matcha_only: bool = False,
 ):
     global logger_handler
     paths = get_path(model_name)
@@ -70,7 +71,7 @@ def initialize(
     logger_handler = logger.add(paths.dataset_path / file_name)
 
     logger.info(
-        f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, freeze_ZH_bert: {freeze_ZH_bert}, freeze_JP_bert: {freeze_JP_bert}, freeze_EN_bert: {freeze_EN_bert}, freeze_style: {freeze_style}, freeze_decoder: {freeze_decoder}, use_jp_extra: {use_jp_extra}"
+        f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, freeze_ZH_bert: {freeze_ZH_bert}, freeze_JP_bert: {freeze_JP_bert}, freeze_EN_bert: {freeze_EN_bert}, freeze_style: {freeze_style}, freeze_decoder: {freeze_decoder}, use_jp_extra: {use_jp_extra}, matcha_only: {matcha_only}"
     )
 
     default_config_path = (
@@ -86,6 +87,10 @@ def initialize(
     config["train"]["epochs"] = epochs
     config["train"]["eval_interval"] = save_every_steps
     config["train"]["log_interval"] = log_interval
+    config["train"]["matcha_only"] = matcha_only
+    if matcha_only:
+        config["model"]["use_matcha"] = True
+        config["model"]["matcha_use_diff_attention"] = True
 
     config["train"]["freeze_EN_bert"] = freeze_EN_bert
     config["train"]["freeze_JP_bert"] = freeze_JP_bert
@@ -275,6 +280,7 @@ def preprocess_all(
     val_per_lang: int,
     log_interval: int,
     yomi_error: str,
+    matcha_only: bool = False,
 ):
     if model_name == "":
         return False, "Error: モデル名を入力してください"
@@ -290,6 +296,7 @@ def preprocess_all(
         freeze_decoder=freeze_decoder,
         use_jp_extra=use_jp_extra,
         log_interval=log_interval,
+        matcha_only=matcha_only,
     )
     if not success:
         return False, message