scripts/build_pr_dataset.py

"""Build data/prs.jsonl from oumi-ai/oumi git history.

Mines squash-merge commits (subject ending in "(#NNNN)"), derives a subsystem
area label from the files each PR changed, and (best-effort) fetches PR bodies
and merge dates via the GitHub GraphQL API using `gh` auth.

Usage:
    git clone --filter=blob:none --single-branch https://github.com/oumi-ai/oumi /tmp/oumi-mine
    uv run python scripts/build_pr_dataset.py --repo /tmp/oumi-mine [--no-bodies]
"""

import argparse
import json
import re
import subprocess
from collections import Counter
from pathlib import Path

OUT_PATH = Path(__file__).resolve().parent.parent / "data" / "prs.jsonl"

AREAS = [
    "training", "inference", "data", "evaluation", "configs",
    "launcher", "cli", "docs", "infra", "other",
]

# Ordered (prefix-or-regex, area) rules; first match wins. Paths not matching
# any rule fall through to "other". Tests and barrel files don't vote.
IGNORE = re.compile(r"^tests/|(^|/)__init__\.py$")
RULES: list[tuple[re.Pattern, str]] = [
    (re.compile(r"^docs/|^notebooks/|\.(md|rst|ipynb)$|^CITATION|^LICENSE"), "docs"),
    (re.compile(r"^\.github/|^scripts/|^pyproject\.toml$|^uv\.lock$|^Makefile$|^Dockerfile|^install\.sh$|^\.pre-commit|^\.(gitignore|dockerignore)|^\.style|^\.readthedocs"), "infra"),
    (re.compile(r"^configs/|^src/(oumi|lema)/core/configs/"), "configs"),
    (re.compile(r"^src/(oumi|lema)/(inference/|infer.*\.py$|deploy/)|^src/(oumi|lema)/core/inference/"), "inference"),
    (re.compile(r"^src/(oumi|lema)/core/(trainers|tuners|callbacks|collators)/|^src/(oumi|lema)/(train|tune)\.py$|^src/(oumi|lema)/performance/|^src/(oumi|lema)/core/distributed\.py$"), "training"),
    (re.compile(r"^src/(oumi|lema)/(evaluation/|evaluate.*\.py$|judges/|judge\.py$|analyze/)|^src/(oumi|lema)/core/(evaluation|analyze)/"), "evaluation"),
    (re.compile(r"^src/(oumi|lema)/(datasets/|synth\.py$)|^src/(oumi|lema)/core/(datasets|tokenizers|processors|feature_generators|synthesis|types)/|^data/"), "data"),
    (re.compile(r"^src/(oumi|lema)/launcher/|^src/(oumi|lema)/core/launcher/"), "launcher"),
    (re.compile(r"^src/(oumi|lema)/cli/"), "cli"),
]

SUBJECT_RE = re.compile(r"^(?P<title>.*)\s+\(#(?P<number>\d+)\)$")
DOMINANCE = 0.70
MAX_BODY_CHARS = 4000


def map_file(path: str) -> str | None:
    if IGNORE.search(path):
        return None
    for pattern, area in RULES:
        if pattern.search(path):
            return area
    return "other"


def derive_area(files: list[str]) -> tuple[str | None, dict[str, int]]:
    """Return (area, votes); area is None when no file dominates."""
    votes = Counter(a for f in files if (a := map_file(f)))
    if not votes:
        return None, {}
    area, top = votes.most_common(1)[0]
    if top / sum(votes.values()) < DOMINANCE:
        return None, dict(votes)
    return area, dict(votes)


def mine_commits(repo: str) -> list[dict]:
    # Single pass: NUL-prefixed header line per commit, then its changed files.
    log = subprocess.run(
        ["git", "-C", repo, "log", "--name-only", "--no-merges",
         "--format=%x00%s", "main"],
        capture_output=True, text=True, check=True,
    ).stdout
    records = []
    for block in log.split("\x00"):
        if not block.strip():
            continue
        subject, _, file_block = block.partition("\n")
        m = SUBJECT_RE.match(subject.strip())
        if not m:
            continue
        files = [ln.strip() for ln in file_block.splitlines() if ln.strip()]
        area, votes = derive_area(files)
        if area is None:
            continue
        records.append({
            "number": int(m["number"]),
            "title": m["title"].strip(),
            "files": files,
            "area": area,
            "area_votes": votes,
        })
    # Dedupe (reverts etc. can repeat a PR number); keep first (newest) occurrence.
    seen: set[int] = set()
    return [r for r in records if not (r["number"] in seen or seen.add(r["number"]))]


def fetch_bodies(records: list[dict], batch_size: int = 50) -> None:
    """Attach `body` and `merged_at` in-place via GitHub GraphQL; skip on failure."""
    for i in range(0, len(records), batch_size):
        batch = records[i:i + batch_size]
        fields = "".join(
            f'p{r["number"]}: pullRequest(number: {r["number"]}) {{ body mergedAt }}\n'
            for r in batch
        )
        query = f'query {{ repository(owner: "oumi-ai", name: "oumi") {{ {fields} }} }}'
        try:
            out = subprocess.run(
                ["gh", "api", "graphql", "-f", f"query={query}"],
                capture_output=True, text=True, check=True, timeout=120,
            ).stdout
            repo_data = json.loads(out)["data"]["repository"]
        except (subprocess.SubprocessError, KeyError, json.JSONDecodeError) as e:
            print(f"  body fetch failed for batch at {i}: {e}; continuing without")
            continue
        for r in batch:
            pr = repo_data.get(f'p{r["number"]}') or {}
            r["body"] = (pr.get("body") or "")[:MAX_BODY_CHARS]
            r["merged_at"] = pr.get("mergedAt")
        print(f"  bodies {i + len(batch)}/{len(records)}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo", default="/tmp/oumi-mine")
    parser.add_argument("--no-bodies", action="store_true")
    args = parser.parse_args()

    records = mine_commits(args.repo)
    print(f"mined {len(records)} PRs with unambiguous areas")
    print("label distribution:", Counter(r["area"] for r in records).most_common())
    if not args.no_bodies:
        fetch_bodies(records)

    records.sort(key=lambda r: -r["number"])
    OUT_PATH.parent.mkdir(exist_ok=True)
    with OUT_PATH.open("w") as f:
        for r in records:
            f.write(json.dumps(r) + "\n")
    print(f"wrote {len(records)} records to {OUT_PATH}")


if __name__ == "__main__":
    main()