| |
| """Probe or download Hugging Face dataset splits for the CyberGym SFT mix. |
| |
| Two modes: |
| |
| Probe (no heavy deps, just stdlib + network) -- inspect a dataset's real |
| schema via the public datasets-server before trusting the manifest: |
| |
| python training/scripts/hf_download.py --probe colin/PrimeVul |
| python training/scripts/hf_download.py --probe-manifest training/configs/datasets.yaml |
| |
| Download (requires `datasets`; runs on the GPU host) -- stream a split to |
| raw JSONL under data/download/<key>/raw.jsonl: |
| |
| python training/scripts/hf_download.py --key primevul |
| python training/scripts/hf_download.py --all --profile pilot |
| python training/scripts/hf_download.py --hf-id colin/PrimeVul --split train --out data/download/primevul/raw.jsonl |
| |
| Auth: set HF_TOKEN (or HUGGINGFACE_HUB_TOKEN) for gated datasets. Probe uses the |
| anonymous datasets-server; gated sets return HTTP 401 there (recorded in the |
| manifest `auth:` field). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import urllib.parse |
| import urllib.request |
| from pathlib import Path |
| from typing import Any |
|
|
| import yaml |
|
|
| DATASETS_SERVER = "https://datasets-server.huggingface.co" |
|
|
|
|
| def read_yaml(path: str | Path) -> dict[str, Any]: |
| with Path(path).open("r", encoding="utf-8") as fh: |
| payload = yaml.safe_load(fh) or {} |
| if not isinstance(payload, dict): |
| raise TypeError(f"Expected a YAML mapping in {path}") |
| return payload |
|
|
|
|
| def _http_get(url: str, timeout: int = 20) -> dict[str, Any]: |
| req = urllib.request.Request(url, headers={"User-Agent": "infosec-hf-probe"}) |
| token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") |
| if token: |
| req.add_header("Authorization", f"Bearer {token}") |
| with urllib.request.urlopen(req, timeout=timeout) as resp: |
| return json.loads(resp.read()) |
|
|
|
|
| def probe(hf_id: str) -> dict[str, Any]: |
| """Return {splits, features, sample_row} or {error} for a dataset.""" |
| enc = urllib.parse.quote(hf_id) |
| try: |
| splits_payload = _http_get(f"{DATASETS_SERVER}/splits?dataset={enc}") |
| except Exception as exc: |
| return {"hf_id": hf_id, "error": repr(exc)} |
| splits = splits_payload.get("splits", []) |
| if not splits: |
| return {"hf_id": hf_id, "error": "no splits", "raw": splits_payload} |
| first = splits[0] |
| cfg, sp = first["config"], first["split"] |
| try: |
| rows_payload = _http_get( |
| f"{DATASETS_SERVER}/first-rows?dataset={enc}" |
| f"&config={urllib.parse.quote(cfg)}&split={urllib.parse.quote(sp)}" |
| ) |
| except Exception as exc: |
| return {"hf_id": hf_id, "splits": splits, "error": f"first-rows: {exc!r}"} |
| features = [ |
| {"name": f["name"], "type": f["type"].get("dtype", f["type"].get("_type"))} |
| for f in rows_payload.get("features", []) |
| ] |
| sample = rows_payload.get("rows", [{}])[0].get("row", {}) if rows_payload.get("rows") else {} |
| return { |
| "hf_id": hf_id, |
| "splits": [{"config": s["config"], "split": s["split"]} for s in splits], |
| "config_used": cfg, |
| "split_used": sp, |
| "features": features, |
| "sample_row": sample, |
| } |
|
|
|
|
| def print_probe(result: dict[str, Any]) -> None: |
| print("=" * 72) |
| print(result["hf_id"]) |
| if "error" in result and "features" not in result: |
| print(f" ERROR: {result['error']}") |
| return |
| print(f" splits: {result.get('splits')}") |
| print(f" used: {result.get('config_used')}/{result.get('split_used')}") |
| print(" features:") |
| for feat in result.get("features", []): |
| print(f" - {feat['name']}: {feat['type']}") |
|
|
|
|
| def resolve_sources(manifest: dict[str, Any]) -> dict[str, dict[str, Any]]: |
| return manifest.get("sources", {}) |
|
|
|
|
| def download_split( |
| hf_id: str, |
| config: str | None, |
| split: str, |
| out_path: Path, |
| max_rows: int | None, |
| ) -> dict[str, Any]: |
| try: |
| from datasets import load_dataset |
| except Exception as exc: |
| raise RuntimeError( |
| "The `datasets` package is required for download mode. " |
| "Install it on the training host (see requirements-cu126.txt)." |
| ) from exc |
|
|
| out_path.parent.mkdir(parents=True, exist_ok=True) |
| token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") |
| stream = load_dataset(hf_id, config or None, split=split, streaming=True, token=token) |
|
|
| written = 0 |
| with out_path.open("w", encoding="utf-8") as out: |
| for row in stream: |
| out.write(json.dumps(row, ensure_ascii=False, default=str) + "\n") |
| written += 1 |
| if max_rows is not None and written >= max_rows: |
| break |
| return {"hf_id": hf_id, "split": split, "out": str(out_path), "rows": written} |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--probe", metavar="HF_ID", help="Probe one dataset schema via the datasets-server.") |
| parser.add_argument("--probe-manifest", metavar="YAML", help="Probe every source in a manifest.") |
| parser.add_argument("--manifest", default="training/configs/datasets.yaml") |
| parser.add_argument("--key", help="Manifest source key to download.") |
| parser.add_argument("--all", action="store_true", help="Download every enabled manifest source.") |
| parser.add_argument("--eval", action="store_true", help="Download the eval_split instead of the train_split.") |
| parser.add_argument("--profile", choices=["full", "pilot"], default="full") |
| parser.add_argument("--hf-id", help="Ad-hoc download: dataset id (bypasses the manifest).") |
| parser.add_argument("--config", help="Ad-hoc download: config name.") |
| parser.add_argument("--split", help="Ad-hoc download: split name.") |
| parser.add_argument("--out", help="Ad-hoc download: output JSONL path.") |
| parser.add_argument("--max-rows", type=int, default=None) |
| parser.add_argument("--include-disabled", action="store_true", help="With --all, also fetch disabled sources.") |
| return parser.parse_args() |
|
|
|
|
| def _cap_for(source: dict[str, Any], defaults: dict[str, Any], profile: str) -> int | None: |
| if profile == "pilot": |
| return source.get("pilot_sample_cap", defaults.get("pilot_sample_cap")) |
| cap = source.get("sample_cap", defaults.get("sample_cap")) |
| return cap |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
|
|
| if args.probe: |
| print_probe(probe(args.probe)) |
| return 0 |
|
|
| if args.probe_manifest: |
| manifest = read_yaml(args.probe_manifest) |
| for key, source in resolve_sources(manifest).items(): |
| print(f"\n### {key} ({source.get('auth', 'public')}, enabled={source.get('enabled')})") |
| print_probe(probe(source["hf_id"])) |
| return 0 |
|
|
| if args.hf_id: |
| out = Path(args.out or f"data/download/_adhoc/{args.hf_id.replace('/', '__')}.jsonl") |
| result = download_split(args.hf_id, args.config, args.split or "train", out, args.max_rows) |
| print(json.dumps(result, indent=2)) |
| return 0 |
|
|
| manifest = read_yaml(args.manifest) |
| defaults = manifest.get("defaults", {}) |
| raw_dir = Path(defaults.get("raw_dir", "data/download")) |
| sources = resolve_sources(manifest) |
|
|
| keys: list[str] |
| if args.all: |
| keys = [ |
| k |
| for k, s in sources.items() |
| if (s.get("enabled", False) or args.include_disabled) |
| ] |
| elif args.key: |
| keys = [args.key] |
| else: |
| print("Nothing to do. Use --probe, --probe-manifest, --key, --all, or --hf-id.", file=sys.stderr) |
| return 2 |
|
|
| results = [] |
| for key in keys: |
| if key not in sources: |
| print(f"Unknown source key: {key}", file=sys.stderr) |
| return 2 |
| source = sources[key] |
| if source.get("auth") == "gated" and not (os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")): |
| print(f"[skip] {key}: gated dataset and no HF_TOKEN set.", file=sys.stderr) |
| continue |
| split_key = "eval_split" if args.eval else "train_split" |
| split = source.get(split_key) |
| if not split: |
| print(f"[skip] {key}: no {split_key} configured.", file=sys.stderr) |
| continue |
| suffix = "eval" if args.eval else "raw" |
| out_path = raw_dir / key / f"{suffix}.jsonl" |
| max_rows = args.max_rows if args.max_rows is not None else _cap_for(source, defaults, args.profile) |
| print(f"[download] {key}: {source['hf_id']} [{split}] -> {out_path} (max_rows={max_rows})") |
| result = download_split(source["hf_id"], source.get("config"), split, out_path, max_rows) |
| result["key"] = key |
| results.append(result) |
| print(json.dumps(result)) |
|
|
| print(json.dumps({"downloaded": results}, indent=2)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
