File size: 2,196 Bytes
d03de05 | 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 | #!/usr/bin/env python3
"""Upload a PP-OCR checkpoint prefix and its training metadata to Hugging Face."""
from __future__ import annotations
import argparse
import json
import os
from datetime import datetime, timezone
from pathlib import Path
from huggingface_hub import HfApi
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--repo-id", required=True)
parser.add_argument("--checkpoint-prefix", type=Path, required=True)
parser.add_argument("--tag", required=True)
parser.add_argument("--metric", type=float, required=True)
parser.add_argument("--config", type=Path, required=True)
parser.add_argument("--dictionary", type=Path, required=True)
parser.add_argument("--token-file", type=Path, default=Path("/root/.hf_token"))
args = parser.parse_args()
token = args.token_file.read_text().strip()
api = HfApi(token=token)
api.create_repo(args.repo_id, repo_type="model", exist_ok=True)
destination = f"checkpoints/{args.tag}"
uploaded = []
for suffix in (".pdparams", ".states"):
source = Path(str(args.checkpoint_prefix) + suffix)
if source.exists():
api.upload_file(
path_or_fileobj=source,
path_in_repo=f"{destination}/{source.name}",
repo_id=args.repo_id,
repo_type="model",
)
uploaded.append(source.name)
for source in (args.config, args.dictionary):
api.upload_file(
path_or_fileobj=source,
path_in_repo=f"{destination}/{source.name}",
repo_id=args.repo_id,
repo_type="model",
)
uploaded.append(source.name)
receipt = {
"tag": args.tag,
"metric": args.metric,
"checkpoint_prefix": args.checkpoint_prefix.name,
"uploaded": uploaded,
"uploaded_at": datetime.now(timezone.utc).isoformat(),
}
api.upload_file(
path_or_fileobj=json.dumps(receipt, indent=2).encode(),
path_in_repo="latest-checkpoint.json",
repo_id=args.repo_id,
repo_type="model",
)
print(json.dumps(receipt))
if __name__ == "__main__":
main()
|