scripts/build_dataset.py

#!/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())