Revert "Refresh graph directly from Hub API"

This reverts commit cfe24a8555f547fd805afa88ed651ba284979b52.

database_refresh/refresh_database.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Build and publish a Neo4j dump from the public Hugging Face Hub API."""
+"""Build and publish a Neo4j dump from the public cfahlgren1/hub-stats dataset."""
 
 from __future__ import annotations
 
@@ -17,7 +17,6 @@ from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
 
 
 SOURCE_REPO = "cfahlgren1/hub-stats"
-HUB_API_SOURCE = "huggingface.co/api"
 DEFAULT_DUMP_REPO = "cnil/genmod-dump-neo4j"
 PARQUET_REVISION = "refs/convert/parquet"
 
@@ -51,106 +50,6 @@ def download_sources(work_dir: Path) -> tuple[Path, Path]:
     return models, datasets
 
 
-def isoformat(value: object) -> str | None:
-    return value.isoformat() if hasattr(value, "isoformat") else None
-
-
-def write_jsonl_snapshot(path: Path, rows: object, kind: str) -> int:
-    count = 0
-    with path.open("w", encoding="utf-8") as output:
-        for info in rows:
-            tags = list(info.tags or [])
-            if kind == "model":
-                safetensors = getattr(info, "safetensors", None)
-                card_data = getattr(info, "card_data", None)
-                base_models = getattr(info, "base_models", None) or {}
-                record = {
-                    "_id": getattr(info, "_id", None),
-                    "id": info.id,
-                    "author": info.author,
-                    "lastModified": isoformat(info.last_modified),
-                    "downloadsAllTime": info.downloads_all_time or info.downloads,
-                    "pipeline_tag": info.pipeline_tag,
-                    "createdAt": isoformat(info.created_at),
-                    "parameters_total": getattr(safetensors, "total", None),
-                    "likes": info.likes,
-                    "license": (
-                        card_data.get("license")
-                        if isinstance(card_data, dict)
-                        else None
-                    ),
-                    "tags": tags,
-                    "base_model_relation": base_models.get("relation"),
-                    "base_model_ids": [
-                        model.get("id")
-                        for model in base_models.get("models", [])
-                        if model.get("id")
-                    ],
-                }
-            else:
-                record = {
-                    "_id": getattr(info, "_id", None),
-                    "id": info.id,
-                    "author": info.author,
-                    "lastModified": isoformat(info.last_modified),
-                    "downloadsAllTime": info.downloads_all_time or info.downloads,
-                    "createdAt": isoformat(info.created_at),
-                    "tags": tags,
-                }
-            output.write(json.dumps(record, ensure_ascii=False) + "\n")
-            count += 1
-            if count % 50000 == 0:
-                log(f"Collected {count:,} {kind} records")
-    log(f"Collected {count:,} {kind} records in total")
-    return count
-
-
-def download_sources_from_api(
-    work_dir: Path,
-    api: HfApi,
-    max_models: int | None,
-    max_datasets: int | None,
-) -> tuple[Path, Path]:
-    """Create deterministic local snapshots without relying on a third-party dataset."""
-    model_fields = [
-        "author",
-        "baseModels",
-        "cardData",
-        "createdAt",
-        "downloads",
-        "downloadsAllTime",
-        "lastModified",
-        "likes",
-        "pipeline_tag",
-        "safetensors",
-        "tags",
-    ]
-    dataset_fields = [
-        "author",
-        "createdAt",
-        "downloads",
-        "downloadsAllTime",
-        "lastModified",
-        "likes",
-        "tags",
-    ]
-    models = work_dir / "models.jsonl"
-    datasets = work_dir / "datasets.jsonl"
-    log("Collecting the model snapshot from the public Hugging Face Hub API")
-    write_jsonl_snapshot(
-        models,
-        api.list_models(expand=model_fields, limit=max_models),
-        "model",
-    )
-    log("Collecting the dataset snapshot from the public Hugging Face Hub API")
-    write_jsonl_snapshot(
-        datasets,
-        api.list_datasets(expand=dataset_fields, limit=max_datasets),
-        "dataset",
-    )
-    return models, datasets
-
-
 def sql_path(path: Path) -> str:
     return str(path).replace("'", "''")
 
@@ -164,47 +63,18 @@ def create_views(
 ) -> None:
     model_limit = f" LIMIT {max_models}" if max_models else ""
     dataset_limit = f" LIMIT {max_datasets}" if max_datasets else ""
-    if models_path.suffix == ".jsonl":
-        model_source = f"read_json_auto('{sql_path(models_path)}', format='newline_delimited')"
-        dataset_source = f"read_json_auto('{sql_path(datasets_path)}', format='newline_delimited')"
-        model_columns = """
-            TRY_CAST(_id AS VARCHAR) AS _id, id, author, lastModified,
-            TRY_CAST(downloadsAllTime AS BIGINT) AS downloadsAllTime,
-            pipeline_tag, createdAt,
-            TRY_CAST(parameters_total AS BIGINT) AS parameters_total,
-            TRY_CAST(likes AS BIGINT) AS likes,
-            TRY_CAST(license AS VARCHAR) AS license, tags,
-            TRY_CAST(base_model_relation AS VARCHAR) AS base_model_relation,
-            TRY_CAST(base_model_ids AS VARCHAR[]) AS base_model_ids
-        """
-        dataset_columns = """
-            TRY_CAST(_id AS VARCHAR) AS _id, id, author, lastModified,
-            TRY_CAST(downloadsAllTime AS BIGINT) AS downloadsAllTime,
-            createdAt, tags
-        """
-    else:
-        model_source = f"read_parquet('{sql_path(models_path)}')"
-        dataset_source = f"read_parquet('{sql_path(datasets_path)}')"
-        model_columns = """
-            _id, id, author, lastModified, downloadsAllTime, pipeline_tag,
-            createdAt, safetensors.total AS parameters_total, likes,
-            json_extract_string(cardData, '$.license') AS license, tags,
-            baseModels.relation AS base_model_relation,
-            list_transform(baseModels.models, model -> model.id) AS base_model_ids
-        """
-        dataset_columns = "_id, id, author, lastModified, downloadsAllTime, createdAt, tags"
     connection.execute(
         f"""
         CREATE VIEW source_models AS
         SELECT * EXCLUDE (_dedupe_rank)
         FROM (
             SELECT
-                {model_columns},
+                *,
                 row_number() OVER (
                     PARTITION BY id
                     ORDER BY lastModified DESC NULLS LAST, _id DESC NULLS LAST
                 ) AS _dedupe_rank
-            FROM {model_source}
+            FROM read_parquet('{sql_path(models_path)}')
             WHERE id IS NOT NULL AND trim(id) <> ''
         )
         WHERE _dedupe_rank = 1
@@ -217,12 +87,12 @@ def create_views(
         SELECT * EXCLUDE (_dedupe_rank)
         FROM (
             SELECT
-                {dataset_columns},
+                *,
                 row_number() OVER (
                     PARTITION BY id
                     ORDER BY lastModified DESC NULLS LAST, _id DESC NULLS LAST
                 ) AS _dedupe_rank
-            FROM {dataset_source}
+            FROM read_parquet('{sql_path(datasets_path)}')
             WHERE id IS NOT NULL AND trim(id) <> ''
         )
         WHERE _dedupe_rank = 1
@@ -233,13 +103,14 @@ def create_views(
         """
         CREATE VIEW base_model_edges AS
         SELECT DISTINCT
-            base_id AS parent_id,
+            base.id AS parent_id,
             child.id AS child_id,
-            COALESCE(child.base_model_relation, 'derived') AS relation_name
+            COALESCE(child.baseModels.relation, 'derived') AS relation_name
         FROM source_models AS child,
-             UNNEST(child.base_model_ids) AS nested(base_id)
-        WHERE base_id IS NOT NULL
-          AND trim(base_id) <> ''
+             UNNEST(child.baseModels.models) AS nested(base)
+        WHERE child.baseModels IS NOT NULL
+          AND base.id IS NOT NULL
+          AND trim(base.id) <> ''
           AND child.id IS NOT NULL
         """
     )
@@ -308,17 +179,17 @@ def prepare_csv_files(
                 pipeline_tag AS task,
                 CAST(createdAt AS VARCHAR) AS created_at,
                 CASE
-                    WHEN parameters_total >= 1000000000
-                        THEN printf('%.1fB', parameters_total / 1000000000.0)
-                    WHEN parameters_total >= 1000000
-                        THEN printf('%.1fM', parameters_total / 1000000.0)
-                    WHEN parameters_total >= 1000
-                        THEN printf('%.1fK', parameters_total / 1000.0)
-                    WHEN parameters_total IS NOT NULL
-                        THEN CAST(parameters_total AS VARCHAR)
+                    WHEN safetensors.total >= 1000000000
+                        THEN printf('%.1fB', safetensors.total / 1000000000.0)
+                    WHEN safetensors.total >= 1000000
+                        THEN printf('%.1fM', safetensors.total / 1000000.0)
+                    WHEN safetensors.total >= 1000
+                        THEN printf('%.1fK', safetensors.total / 1000.0)
+                    WHEN safetensors.total IS NOT NULL
+                        THEN CAST(safetensors.total AS VARCHAR)
                 END AS parameters,
                 likes,
-                license
+                json_extract_string(cardData, '$.license') AS license
             FROM source_models
             WHERE id IS NOT NULL AND trim(id) <> ''
         ),
@@ -468,14 +339,13 @@ def build_dump(files: dict[str, Path], output_dir: Path, neo4j_admin: str) -> Pa
 
 def write_metadata(
     output_dir: Path,
-    source_repo: str,
     source_revision: str,
     model_count: int,
     dataset_count: int,
 ) -> Path:
     metadata = {
         "built_at": datetime.now(timezone.utc).isoformat(),
-        "source_repo": source_repo,
+        "source_repo": SOURCE_REPO,
         "source_revision": source_revision,
         "model_count": model_count,
         "dataset_count": dataset_count,
@@ -487,11 +357,10 @@ def write_metadata(
 
 def parquet_unique_id_count(path: Path) -> int:
     connection = duckdb.connect()
-    reader = "read_json_auto" if path.suffix == ".jsonl" else "read_parquet"
     count = connection.execute(
         f"""
         SELECT count(DISTINCT id)
-        FROM {reader}('{sql_path(path)}')
+        FROM read_parquet('{sql_path(path)}')
         WHERE id IS NOT NULL AND trim(id) <> ''
         """
     ).fetchone()[0]
@@ -528,7 +397,7 @@ def publish_dump(
                 path_or_fileobj=str(metadata_path),
             ),
         ],
-        commit_message="Refresh Neo4j graph from the Hugging Face Hub API",
+        commit_message="Refresh Neo4j graph from cfahlgren1/hub-stats",
     )
 
 
@@ -544,11 +413,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--dump-repo", default=os.getenv("NEO4J_DUMP_REPO", DEFAULT_DUMP_REPO))
     parser.add_argument("--dump-revision", default=os.getenv("NEO4J_DUMP_REVISION", "main"))
     parser.add_argument("--neo4j-admin", default=os.getenv("NEO4J_ADMIN", "neo4j-admin"))
-    parser.add_argument(
-        "--source",
-        choices=("hub-api", "hub-stats"),
-        default=os.getenv("GENMOD_REFRESH_SOURCE", "hub-api"),
-    )
     parser.add_argument("--models-parquet", type=Path)
     parser.add_argument("--datasets-parquet", type=Path)
     parser.add_argument("--max-models", type=int)
@@ -567,25 +431,13 @@ def main() -> int:
     args.work_dir.mkdir(parents=True, exist_ok=True)
 
     api = HfApi()
+    source_revision = api.dataset_info(SOURCE_REPO).sha
     if bool(args.models_parquet) != bool(args.datasets_parquet):
         raise SystemExit("Provide both --models-parquet and --datasets-parquet.")
     if args.models_parquet:
         models_path, datasets_path = args.models_parquet, args.datasets_parquet
-        source_repo = "local-files"
-        source_revision = "provided"
-    elif args.source == "hub-stats":
-        source_repo = SOURCE_REPO
-        source_revision = api.dataset_info(SOURCE_REPO).sha
-        models_path, datasets_path = download_sources(args.work_dir)
     else:
-        source_repo = HUB_API_SOURCE
-        source_revision = datetime.now(timezone.utc).isoformat()
-        models_path, datasets_path = download_sources_from_api(
-            args.work_dir,
-            api,
-            args.max_models,
-            args.max_datasets,
-        )
+        models_path, datasets_path = download_sources(args.work_dir)
 
     csv_dir = args.work_dir / "csv"
     files = prepare_csv_files(
@@ -604,7 +456,6 @@ def main() -> int:
     dataset_count = args.max_datasets or parquet_unique_id_count(datasets_path)
     metadata_path = write_metadata(
         args.work_dir,
-        source_repo,
         source_revision,
         model_count,
         dataset_count,

database_refresh/run_weekly_refresh.sh

@@ -1,6 +1,7 @@
 #!/bin/bash
 set -e
 apt-get update
-apt-get install -y python3-pip
-pip3 install duckdb==1.3.2 huggingface-hub==1.25.1
-exec python3 /workspace/database_refresh/refresh_database.py --source hub-api
+apt-get install -y python3-pip wget
+pip3 install duckdb==1.3.2 huggingface-hub==0.31.4
+wget -qO /tmp/r.py https://huggingface.co/spaces/cnil/genmod/resolve/refresh-hub-database/database_refresh/refresh_database.py
+exec python3 /tmp/r.py