Add training-kit/upload_ppocr_checkpoint.py

training-kit/upload_ppocr_checkpoint.py

@@ -0,0 +1,68 @@
+#!/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()