fix

configure_matcha_ft.py

@@ -0,0 +1,79 @@
+"""Configure a dataset for Matcha-only or joint fine-tuning."""
+
+import argparse
+import json
+import shutil
+from pathlib import Path
+from typing import Optional
+
+
+def configure_matcha_ft(
+    config_path: Path,
+    matcha_only: bool = True,
+    learning_rate: Optional[float] = None,
+    epochs: Optional[int] = None,
+    save_every_steps: Optional[int] = None,
+) -> dict:
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+    with config_path.open(encoding="utf-8") as stream:
+        config = json.load(stream)
+
+    train = config.setdefault("train", {})
+    model = config.setdefault("model", {})
+    train["matcha_only"] = matcha_only
+    model["use_matcha"] = True
+    model["matcha_use_diff_attention"] = True
+    if learning_rate is not None:
+        if learning_rate <= 0:
+            raise ValueError("Learning rate must be positive")
+        train["learning_rate"] = learning_rate
+    if epochs is not None:
+        if epochs < 1:
+            raise ValueError("Epochs must be at least 1")
+        train["epochs"] = epochs
+    if save_every_steps is not None:
+        if save_every_steps < 1:
+            raise ValueError("Save interval must be at least 1")
+        train["eval_interval"] = save_every_steps
+
+    backup_path = config_path.with_suffix(config_path.suffix + ".bak")
+    if not backup_path.exists():
+        shutil.copy2(config_path, backup_path)
+    with config_path.open("w", encoding="utf-8") as stream:
+        json.dump(config, stream, indent=2, ensure_ascii=False)
+        stream.write("\n")
+    return config
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Configure Matcha Flow fine-tuning for an existing dataset"
+    )
+    parser.add_argument("--config", type=Path, required=True)
+    parser.add_argument(
+        "--joint",
+        action="store_true",
+        help="Jointly fine-tune the full model instead of Matcha-only training",
+    )
+    parser.add_argument("--learning-rate", type=float)
+    parser.add_argument("--epochs", type=int)
+    parser.add_argument("--save-every-steps", type=int)
+    args = parser.parse_args()
+
+    config = configure_matcha_ft(
+        config_path=args.config,
+        matcha_only=not args.joint,
+        learning_rate=args.learning_rate,
+        epochs=args.epochs,
+        save_every_steps=args.save_every_steps,
+    )
+    mode = "joint fine-tuning" if args.joint else "Matcha-only fine-tuning"
+    print(f"Configured {args.config} for {mode}")
+    print(f"learning_rate={config['train']['learning_rate']}")
+    print(f"epochs={config['train']['epochs']}")
+    print(f"save_every_steps={config['train']['eval_interval']}")
+
+
+if __name__ == "__main__":
+    main()