#!/usr/bin/env python3 """Build the prebuilt PR dataset for the "New Model Additions" Space. Crawls the GitHub API for merged "new model" PRs in huggingface/transformers and vllm-project/vllm, replicates the vLLM filtering done in the browser (js/utils.js `addsNewModelFile` + data/blacklist.json), and writes: - prs.parquet : one row per kept PR (columns: repo, number, title, html_url, merged_at) - meta.json : generated_at, per-repo cutoff (max merged_at) and counts Both files are uploaded to the HF dataset repo so the browser can read the bulk history instead of crawling GitHub on every visit (it only tops up newer PRs). Usage: GITHUB_TOKEN=ghp_xxx HF_TOKEN=hf_xxx python scripts/build_dataset.py # add --dry-run to skip the upload and just write the files locally The script is run manually by the maintainer; there is no CI for it. """ from __future__ import annotations import argparse import json import os import sys import time from datetime import datetime, timezone from pathlib import Path import pandas as pd import requests # ── Config (mirrors js/constants.js — keep in sync) ───────────────────────── REPOS = { "hf": {"owner": "huggingface", "repo": "transformers", "label": "New model"}, "vllm": {"owner": "vllm-project", "repo": "vllm", "label": "new-model"}, } VLLM_MODELS_PATH = "vllm/model_executor/models/" VLLM_INFRA_FILES = { "adapters.py", "config.py", "interfaces_base.py", "interfaces.py", "module_mapping.py", "registry.py", } DATASET_REPO = "hmellor/new-model-additions" ROOT = Path(__file__).resolve().parent.parent BLACKLIST_PATH = ROOT / "data" / "blacklist.json" GITHUB_API = "https://api.github.com" SEARCH_URL = f"{GITHUB_API}/search/issues" PER_PAGE = 100 SEARCH_CAP = 1000 # GitHub Search API hard limit per query # ── GitHub helpers ────────────────────────────────────────────────────────── def make_session(token: str | None) -> requests.Session: s = requests.Session() s.headers.update({ "Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28", }) if token: s.headers.update({"Authorization": f"Bearer {token}"}) return s def gh_get(session: requests.Session, url: str, params: dict | None = None) -> requests.Response: """GET with simple rate-limit backoff. On non-rate-limit errors, surface the GitHub response body (422s carry the exact reason in message/errors).""" while True: r = session.get(url, params=params) if r.status_code == 403 and r.headers.get("X-RateLimit-Remaining") == "0": reset = int(r.headers.get("X-RateLimit-Reset", "0")) wait = max(reset - int(time.time()), 1) + 1 print(f" rate limited; sleeping {wait}s…", file=sys.stderr) time.sleep(wait) continue if not r.ok: try: body = r.json() except Exception: body = r.text print(f"\nGitHub {r.status_code} for {r.url}\n body: {body}", file=sys.stderr) r.raise_for_status() return r def search_all_merged(session: requests.Session, cfg: dict) -> list[dict]: """All merged PRs with the repo's label, working around the 1000-result cap by windowing on `created` when needed. Dedupes by PR number.""" base_q = ( f'repo:{cfg["owner"]}/{cfg["repo"]}' f' label:"{cfg["label"]}" is:pr is:merged' ) items: dict[int, dict] = {} cursor: str | None = None prev_cursor: str | None = None while True: # Sort via the query qualifier, not the deprecated sort/order URL params: # the new advanced-search backend (used for authenticated requests) 422s # on sort/order params. advanced_search=true is the supported path now # that GitHub is sunsetting the legacy issue-search backend. q = base_q + (f" created:>={cursor}" if cursor else "") + " sort:created-asc" window_total = None last_created = None for page in range(1, SEARCH_CAP // PER_PAGE + 1): # up to 10 pages = 1000 params = {"q": q, "per_page": PER_PAGE, "page": page, "advanced_search": "true"} data = gh_get(session, SEARCH_URL, params).json() window_total = data.get("total_count", 0) batch = data.get("items", []) if not batch: break for it in batch: items[it["number"]] = it last_created = it.get("created_at") print(f" [{cfg['repo']}] {len(items)} PRs collected " f"(window total {window_total})", end="\r", file=sys.stderr) if window_total is None or window_total <= SEARCH_CAP: break if not last_created or last_created == prev_cursor: print(f"\n WARNING: cannot advance past {last_created}; " f"history may be truncated.", file=sys.stderr) break prev_cursor, cursor = cursor, last_created print(file=sys.stderr) return list(items.values()) def adds_new_model_file(files: list[dict]) -> bool: """Mirror of js/utils.js:addsNewModelFile.""" for f in files: if f.get("status") != "added": continue name = f.get("filename", "") if not name.startswith(VLLM_MODELS_PATH): continue rel = name[len(VLLM_MODELS_PATH):] if rel not in VLLM_INFRA_FILES and not rel.startswith("transformers/"): return True return False def pr_adds_model(session: requests.Session, cfg: dict, number: int) -> bool: """Fetch the PR's changed files (paginated) and apply the predicate.""" page = 1 while True: url = f"{GITHUB_API}/repos/{cfg['owner']}/{cfg['repo']}/pulls/{number}/files" files = gh_get(session, url, {"per_page": PER_PAGE, "page": page}).json() if adds_new_model_file(files): return True if len(files) < PER_PAGE: return False page += 1 def merged_at(pr: dict) -> str | None: return (pr.get("pull_request") or {}).get("merged_at") or pr.get("closed_at") # ── Build ─────────────────────────────────────────────────────────────────── def build_rows(search_session: requests.Session, api_session: requests.Session, blacklist: set[int], check_files: bool = True) -> tuple[list[dict], list[int]]: rows: list[dict] = [] rejected: list[int] = [] for key, cfg in REPOS.items(): print(f"Fetching {cfg['owner']}/{cfg['repo']} (label '{cfg['label']}')…") prs = search_all_merged(search_session, cfg) print(f" {len(prs)} merged PRs with the label") if key == "vllm": kept = [] for i, pr in enumerate(prs, 1): n = pr["number"] if n in blacklist: rejected.append(n) continue if not check_files: # Blacklist-only mode: trust the blacklist, skip per-PR file # checks (avoids the 60/hr anonymous REST limit). May let a few # recent infra-only PRs through until a tokened rebuild. kept.append(pr) continue if pr_adds_model(api_session, cfg, n): kept.append(pr) else: rejected.append(n) if i % 25 == 0: print(f" checked files for {i}/{len(prs)} PRs", end="\r") note = "" if check_files else " (blacklist only — file checks skipped)" print(f"\n vLLM kept {len(kept)} after filtering " f"({len(prs) - len(kept)} rejected){note}") prs = kept for pr in prs: rows.append({ "repo": key, "number": int(pr["number"]), "title": pr.get("title") or "", "html_url": pr.get("html_url") or "", "merged_at": merged_at(pr), }) return rows, rejected def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--dry-run", action="store_true", help="write files locally but skip the HF upload") ap.add_argument("--out", default=str(ROOT / "build"), help="local output directory") ap.add_argument("--no-file-check", action="store_true", help="skip the per-PR vLLM file checks and trust data/blacklist.json " "(fast, anonymous-friendly; may admit a few recent infra-only PRs)") args = ap.parse_args() gh_token = os.environ.get("GITHUB_TOKEN") if not gh_token and not args.no_file_check: print("WARNING: GITHUB_TOKEN not set — the vLLM file checks will hit the 60/hr " "anonymous rate limit (~3h with backoff). Use --no-file-check for a fast " "blacklist-only first build.", file=sys.stderr) blacklist = set(json.loads(BLACKLIST_PATH.read_text())) # The Search API ignores a fine-grained token's repo scoping and 422s on # public repos the token isn't granted ("…do not have permission to view"), # so search is always unauthenticated (it's only a few cheap calls). The # token is used for the high-volume per-PR file checks, which read public # data fine and benefit from the higher authenticated rate limit. search_session = make_session(None) api_session = make_session(gh_token) rows, rejected = build_rows(search_session, api_session, blacklist, check_files=not args.no_file_check) df = pd.DataFrame(rows, columns=["repo", "number", "title", "html_url", "merged_at"]) df["number"] = df["number"].astype("int64") # Per-repo cutoff = latest merged_at present in the dataset cutoff, counts = {}, {} for key in REPOS: sub = df[df["repo"] == key]["merged_at"].dropna() cutoff[key] = max(sub) if len(sub) else None counts[key] = int((df["repo"] == key).sum()) meta = { "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), "cutoff": cutoff, "counts": counts, "schema_version": 1, } out = Path(args.out) out.mkdir(parents=True, exist_ok=True) parquet_path = out / "prs.parquet" meta_path = out / "meta.json" df.to_parquet(parquet_path, index=False) meta_path.write_text(json.dumps(meta, indent=2) + "\n") print(f"\nWrote {parquet_path} ({counts}) and {meta_path}") print(f"cutoff: {cutoff}") # Newly-rejected vLLM PRs not already in the blacklist — paste into data/blacklist.json new_rejects = sorted(set(rejected) - blacklist) if new_rejects: print(f"\n{len(new_rejects)} rejected vLLM PR(s) NOT in data/blacklist.json " f"(consider adding them):") print(", ".join(str(n) for n in new_rejects)) if args.dry_run: print("\n--dry-run: skipping upload.") return 0 from huggingface_hub import HfApi hf_token = os.environ.get("HF_TOKEN") api = HfApi(token=hf_token) api.create_repo(repo_id=DATASET_REPO, repo_type="dataset", exist_ok=True) for path in (parquet_path, meta_path): api.upload_file( path_or_fileobj=str(path), path_in_repo=path.name, repo_id=DATASET_REPO, repo_type="dataset", commit_message=f"Update {path.name} ({meta['generated_at']})", ) print(f"Uploaded {path.name} → datasets/{DATASET_REPO}") return 0 if __name__ == "__main__": raise SystemExit(main())