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