scripts/publish_catalogs.py

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from app.catalog import Catalog  # noqa: E402
from app.datasets import DatasetCatalog  # noqa: E402


def validate_catalogs() -> tuple[Path, Path]:
    runtime = Catalog.model_validate_json((ROOT / "data" / "index.json").read_text())
    community_path = ROOT / "data" / "community-models.json"
    community = Catalog.model_validate_json(community_path.read_text())
    dataset_path = ROOT / "data" / "datasets.json"
    DatasetCatalog.model_validate_json(dataset_path.read_text())

    runtime_slugs = {model.slug for model in runtime.models}
    community_slugs = {model.slug for model in community.models}
    if overlap := runtime_slugs & community_slugs:
        raise ValueError(f"runtime and community catalogs overlap: {sorted(overlap)}")
    runtime_repos = {model.source.repo_id.casefold() for model in runtime.models}
    community_repos = {model.source.repo_id.casefold() for model in community.models}
    if overlap := runtime_repos & community_repos:
        raise ValueError(f"runtime and community repositories overlap: {sorted(overlap)}")
    if any(model.verification.status == "verified" for model in community.models):
        raise ValueError("community models cannot be marked verified")
    return community_path, dataset_path


def main() -> int:
    parser = argparse.ArgumentParser(description="publish AirGPT Hub community catalogs")
    parser.add_argument("--repo-id", default="airgpt/embodied-model-index")
    parser.add_argument("--publish", action="store_true")
    args = parser.parse_args()
    community_path, dataset_path = validate_catalogs()
    print(f"validated {community_path.name} and {dataset_path.name}")
    if not args.publish:
        print("Preview only. Re-run with --publish after review.")
        return 0

    token = os.getenv("HF_TOKEN")
    if not token:
        parser.error("HF_TOKEN is required with --publish")
    try:
        from huggingface_hub import CommitOperationAdd, HfApi
    except ImportError as exc:
        raise SystemExit("install huggingface_hub before publishing") from exc
    result = HfApi(token=token).create_commit(
        repo_id=args.repo_id,
        repo_type="dataset",
        operations=[
            CommitOperationAdd(
                path_in_repo="community-models.json",
                path_or_fileobj=community_path,
            ),
            CommitOperationAdd(
                path_in_repo="datasets.json",
                path_or_fileobj=dataset_path,
            ),
        ],
        commit_message="Expand AirGPT Hub model and dataset catalogs",
    )
    print(result.commit_url)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())