File size: 1,920 Bytes
be90022 |
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 |
#!/usr/bin/env python3
"""
HuggingFace Checkpoint Uploader for AGILLM-3
Usage: python hf_upload.py <checkpoint.pt> [--message "commit msg"]
Requires: HF_TOKEN env var or ~/.huggingface/token
"""
import os
import sys
import argparse
from pathlib import Path
from datetime import datetime
try:
from huggingface_hub import HfApi, upload_file
except ImportError:
print("Installing huggingface_hub...")
os.system("pip install huggingface_hub -q")
from huggingface_hub import HfApi, upload_file
REPO_ID = "OpenTransformer/AGILLM-3-large"
def upload_checkpoint(ckpt_path: str, message: str = None):
path = Path(ckpt_path)
if not path.exists():
print(f"Error: {ckpt_path} not found")
return False
token = os.environ.get("HF_TOKEN") or None
api = HfApi(token=token)
# Extract step from filename
step = "unknown"
if "step" in path.stem:
step = path.stem.split("step")[-1].split(".")[0]
ts = datetime.now().strftime("%Y-%m-%d %H:%M")
commit_msg = message or f"Checkpoint step {step} - {ts}"
print(f"Uploading {path.name} to {REPO_ID}...")
print(f"Size: {path.stat().st_size / 1e9:.2f} GB")
try:
upload_file(
path_or_fileobj=str(path),
path_in_repo=f"checkpoints/{path.name}",
repo_id=REPO_ID,
commit_message=commit_msg,
token=token
)
print(f"SUCCESS: https://huggingface.co/{REPO_ID}")
return True
except Exception as e:
print(f"Upload failed: {e}")
return False
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("checkpoint", help="Path to checkpoint .pt file")
parser.add_argument("--message", "-m", help="Commit message")
args = parser.parse_args()
success = upload_checkpoint(args.checkpoint, args.message)
sys.exit(0 if success else 1)
|