80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
"""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()
|