File size: 2,762 Bytes
4cd3cad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """
Wrapper around sft_12hz.py that pushes each checkpoint to HF Hub after every epoch.
Reads sft_12hz.py's checkpoint output dirs and uploads them.
"""
import argparse
import subprocess
import sys
import os
import shutil
from pathlib import Path
from huggingface_hub import HfApi, create_repo
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--init_model_path", required=True)
p.add_argument("--output_model_path", required=True)
p.add_argument("--train_jsonl", required=True)
p.add_argument("--batch_size", type=int, default=4)
p.add_argument("--lr", type=float, default=2e-6)
p.add_argument("--num_epochs", type=int, default=5)
p.add_argument("--speaker_name", default="mostafa_speaker")
p.add_argument("--hf_repo", required=True)
return p.parse_args()
def push_checkpoint(local_dir: Path, hf_repo: str, epoch: int):
api = HfApi()
try:
create_repo(hf_repo, exist_ok=True, repo_type="model")
except Exception:
pass
print(f"\n→ Pushing checkpoint-epoch-{epoch} to {hf_repo} ...")
api.upload_folder(
folder_path=str(local_dir),
repo_id=hf_repo,
repo_type="model",
commit_message=f"epoch {epoch} checkpoint",
)
print(f"✓ Pushed checkpoint-epoch-{epoch}")
def main():
args = parse_args()
output_dir = Path(args.output_model_path)
# Run one epoch at a time so we can push after each
for epoch in range(args.num_epochs):
print(f"\n{'='*50}")
print(f"Training epoch {epoch+1}/{args.num_epochs}")
print(f"{'='*50}")
cmd = [
sys.executable, "sft_12hz.py",
"--init_model_path", args.init_model_path,
"--output_model_path", args.output_model_path,
"--train_jsonl", args.train_jsonl,
"--batch_size", str(args.batch_size),
"--lr", str(args.lr),
"--num_epochs", str(epoch + 1), # train up to this epoch
"--speaker_name", args.speaker_name,
]
result = subprocess.run(cmd, check=True)
# Push the latest checkpoint
ckpt_dir = output_dir / f"checkpoint-epoch-{epoch}"
if ckpt_dir.exists():
push_checkpoint(ckpt_dir, args.hf_repo, epoch)
# --- NEW CLEANUP BLOCK ---
try:
shutil.rmtree(ckpt_dir)
print(f"✓ Deleted local folder {ckpt_dir} to free up SSD space.")
except Exception as e:
print(f"Warning: Could not delete {ckpt_dir}: {e}")
# -------------------------
else:
print(f"Warning: {ckpt_dir} not found, skipping push")
print("\nAll epochs done and pushed.")
if __name__ == "__main__":
main() |