|
|
|
|
|
""" |
|
|
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) |
|
|
|
|
|
|
|
|
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) |
|
|
|
