scripts/reconstruct_external_data.py

#!/usr/bin/env python3
"""
Reconstruct external (augmented) datasets for the benchmark.

This is a convenience wrapper around `run_all_adapters.py` that:
  - Optionally clones public (redistributable) source repos into a local folder
  - Runs Core adapters (redistributable) with the correct file paths
  - Runs Extended adapters (research-only) individually, with clear guidance if access is blocked

IMPORTANT LICENSING NOTE:
  - Core outputs (MIT/BSD/Apache) may be redistributed with the benchmark.
  - Extended outputs (e.g., CC-BY-NC, research clauses) must NOT be redistributed.
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path
from typing import Optional, List


REPOS_CORE = {
    "HarmBench": "https://github.com/centerforaisafety/HarmBench.git",
    "InjecAgent": "https://github.com/uiuc-kang-lab/InjecAgent.git",
    "ToolEmu": "https://github.com/ryoungj/ToolEmu.git",
    "FigStep": "https://github.com/ThuCCSLab/FigStep.git",
}


def run(cmd: List[str], cwd: Path) -> int:
    print("\n$ " + " ".join(cmd))
    return subprocess.call(cmd, cwd=str(cwd))


def ensure_git_available() -> None:
    try:
        subprocess.check_call(["git", "--version"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception as e:
        raise RuntimeError("git is required for --clone-core-sources but was not found") from e


def clone_repo(url: str, dest: Path) -> None:
    if dest.exists():
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    subprocess.check_call(["git", "clone", "--depth", "1", url, str(dest)])


def main() -> int:
    scripts_dir = Path(__file__).resolve().parent
    repo_root = scripts_dir.parent
    run_all = scripts_dir / "run_all_adapters.py"

    parser = argparse.ArgumentParser(description="Reconstruct external augmented datasets")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=repo_root / "data" / "external_augmented",
        help="Base output directory (default: <repo>/data/external_augmented)",
    )
    parser.add_argument(
        "--sources-dir",
        type=Path,
        default=repo_root / "external_sources",
        help="Where to place cloned/downloaded external sources (default: <repo>/external_sources)",
    )
    parser.add_argument(
        "--core",
        action="store_true",
        help="Run Core (redistributable) adapters",
    )
    parser.add_argument(
        "--extended",
        action="store_true",
        help="Run Extended (research-only) adapters (do not redistribute outputs)",
    )
    parser.add_argument(
        "--clone-core-sources",
        action="store_true",
        help="Clone Core source repos (HarmBench/InjecAgent/ToolEmu/FigStep) into --sources-dir",
    )
    parser.add_argument(
        "--figstep-copy-images",
        action="store_true",
        help="Copy FigStep images into output dir (Core; increases size)",
    )
    parser.add_argument(
        "--tensortrust-attacks-jsonl-path",
        type=Path,
        default=None,
        help="Optional: path to TensorTrust exported attacks JSONL (raw_dump_attacks.jsonl).",
    )
    parser.add_argument(
        "--mmsafetybench-configs",
        type=str,
        default=None,
        help="Comma-separated MM-SafetyBench configs (Extended; optional)",
    )
    parser.add_argument(
        "--mmsafetybench-splits",
        type=str,
        default=None,
        help="Comma-separated MM-SafetyBench splits (Extended; optional)",
    )
    parser.add_argument(
        "--mmsafetybench-include-text-only",
        action="store_true",
        help="Include MM-SafetyBench Text_only split (Extended; optional)",
    )
    parser.add_argument(
        "--mmsafetybench-export-images",
        action="store_true",
        help="Export MM-SafetyBench images into output dir (Extended; DO NOT REDISTRIBUTE)",
    )

    args = parser.parse_args()

    if not args.core and not args.extended:
        print("Error: at least one of --core or --extended is required")
        return 2

    output_dir = args.output_dir.resolve()
    sources_dir = args.sources_dir.resolve()

    harmbench_csv = sources_dir / "HarmBench" / "data" / "behavior_datasets" / "harmbench_behaviors_text_all.csv"
    injecagent_data = sources_dir / "InjecAgent" / "data"
    toolemu_repo = sources_dir / "ToolEmu"
    figstep_repo = sources_dir / "FigStep"

    if args.clone_core_sources and args.core:
        ensure_git_available()
        print(f"Cloning Core source repos into {sources_dir} ...")
        for name, url in REPOS_CORE.items():
            dest = sources_dir / name
            print(f"  - {name}: {url} -> {dest}")
            clone_repo(url, dest)

    if args.core:
        if not harmbench_csv.exists():
            print(f"\nMissing HarmBench CSV: {harmbench_csv}")
            print("Either clone via --clone-core-sources or provide the file manually under --sources-dir.")
            return 2
        if not injecagent_data.exists():
            print(f"\nMissing InjecAgent data dir: {injecagent_data}")
            print("Either clone via --clone-core-sources or provide the repo manually under --sources-dir.")
            return 2
        if not (toolemu_repo / "assets" / "all_cases.json").exists():
            print(f"\nMissing ToolEmu assets/all_cases.json under: {toolemu_repo}")
            print("Either clone via --clone-core-sources or provide the repo manually under --sources-dir.")
            return 2
        if not (figstep_repo / "data" / "question" / "safebench.csv").exists():
            print(f"\nMissing FigStep data/question/safebench.csv under: {figstep_repo}")
            print("Either clone via --clone-core-sources or provide the repo manually under --sources-dir.")
            return 2

        cmd = [
            sys.executable,
            str(run_all),
            "--core-only",
            "--output-dir",
            str(output_dir),
            "--harmbench-csv-path",
            str(harmbench_csv),
            "--injecagent-data-dir",
            str(injecagent_data),
            "--toolemu-repo-dir",
            str(toolemu_repo),
            "--figstep-repo-dir",
            str(figstep_repo),
        ]
        if args.figstep_copy_images:
            cmd.append("--figstep-copy-images")
        if args.tensortrust_attacks_jsonl_path:
            cmd.extend(["--tensortrust-attacks-jsonl-path", str(args.tensortrust_attacks_jsonl_path)])

        rc = run(cmd, cwd=scripts_dir)
        if rc != 0:
            return rc

    if args.extended:
        print("\n" + "=" * 70)
        print("EXTENDED DATASETS (RESEARCH ONLY) — DO NOT REDISTRIBUTE OUTPUTS")
        print("=" * 70)
        print("If Mindgard fails with 'gated/403', request access on HuggingFace and run:")
        print("  huggingface-cli login")
        print("=" * 70 + "\n")

        # Run each extended adapter separately so one failure doesn't block the rest.
        extended_adapters = ["agentharm", "mindgard", "mmsafetybench"]

        mmsb_args: List[str] = []
        if args.mmsafetybench_configs:
            mmsb_args.extend(["--mmsafetybench-configs", args.mmsafetybench_configs])
        if args.mmsafetybench_splits:
            mmsb_args.extend(["--mmsafetybench-splits", args.mmsafetybench_splits])
        if args.mmsafetybench_include_text_only:
            mmsb_args.append("--mmsafetybench-include-text-only")
        if args.mmsafetybench_export_images:
            mmsb_args.append("--mmsafetybench-export-images")

        overall_ok = True
        for adapter in extended_adapters:
            cmd = [
                sys.executable,
                str(run_all),
                "--adapter",
                adapter,
                "--output-dir",
                str(output_dir),
            ]
            if adapter == "mmsafetybench":
                cmd.extend(mmsb_args)

            rc = run(cmd, cwd=scripts_dir)
            if rc != 0:
                overall_ok = False
                print(f"\n[WARN] Adapter '{adapter}' failed (exit={rc}). Continuing.\n")

        if not overall_ok:
            print("One or more Extended adapters failed. This is often expected (e.g., gated access).")
            return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main())