guidedbench 0.1.0__py3-none-any.whl

guidedbench/__init__.py

@@ -0,0 +1,27 @@
+from ._version import BENCHMARK_VERSION, __version__
+from .dataset import (
+    DATASET_REPO_ID,
+    DATASET_REVISION,
+    export_jsonl,
+    get_case,
+    load_cases,
+)
+from .evaluator import GuidedBenchEvaluator
+from .models import EvaluationResult, GuidedBenchCase, Guideline, ScoringPointResult
+from .parsing import EvaluationParseError
+
+__all__ = [
+    "BENCHMARK_VERSION",
+    "DATASET_REPO_ID",
+    "DATASET_REVISION",
+    "EvaluationParseError",
+    "EvaluationResult",
+    "GuidedBenchCase",
+    "GuidedBenchEvaluator",
+    "Guideline",
+    "ScoringPointResult",
+    "__version__",
+    "export_jsonl",
+    "get_case",
+    "load_cases",
+]

guidedbench/_version.py

@@ -0,0 +1,2 @@
+__version__ = "0.1.0"
+BENCHMARK_VERSION = "1.0"

guidedbench/cli.py

@@ -0,0 +1,423 @@
+from __future__ import annotations
+
+import argparse
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+import hashlib
+import json
+import os
+import sys
+import threading
+import uuid
+from pathlib import Path
+from typing import Any, Iterable
+
+from ._version import BENCHMARK_VERSION, __version__
+from .dataset import export_jsonl, get_case, load_cases
+from .evaluator import GuidedBenchEvaluator
+from .judges import (
+    AnthropicJudge,
+    MissingDependencyError,
+    OpenAICompatibleJudge,
+    OpenAIJudge,
+    TransformersJudge,
+)
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="guidedbench",
+        description="Run the GuidedBench guideline-grounded jailbreak evaluation.",
+    )
+    parser.add_argument("--version", action="version", version=f"guidedbench {__version__}")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    subparsers.add_parser("info", help="Show benchmark version and split counts.")
+
+    export = subparsers.add_parser(
+        "export", help="Export the benchmark from Hugging Face as JSON Lines."
+    )
+    export.add_argument("output", type=Path)
+    export.add_argument("--subset", choices=("all", "core", "additional"), default="all")
+    export.add_argument("--topic")
+
+    evaluate = subparsers.add_parser("evaluate", help="Score victim-model responses from JSON Lines.")
+    evaluate.add_argument("input", type=Path, help="JSONL with id/index/question and response fields.")
+    evaluate.add_argument("--output", type=Path, required=True)
+    evaluate.add_argument(
+        "--provider",
+        required=True,
+        choices=("openai", "openai-compatible", "anthropic", "transformers"),
+    )
+    evaluate.add_argument("--model", required=True)
+    evaluate.add_argument("--base-url", help="Base URL for an OpenAI-compatible endpoint.")
+    evaluate.add_argument(
+        "--api-key-env",
+        help="Environment variable containing the provider API key.",
+    )
+    evaluate.add_argument(
+        "--api-style",
+        choices=("chat_completions", "responses"),
+        default="chat_completions",
+        help="Protocol for an OpenAI-compatible endpoint.",
+    )
+    evaluate.add_argument("--max-tokens", type=int, default=800)
+    evaluate.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help=(
+            "Concurrent remote-judge requests. Every completed evaluation is "
+            "persisted as an independent atomic checkpoint before final gather."
+        ),
+    )
+    evaluate.add_argument(
+        "--checkpoint-dir",
+        type=Path,
+        help="Checkpoint directory (default: <output>.parts).",
+    )
+    evaluate.add_argument(
+        "--keep-checkpoints",
+        action="store_true",
+        help="Keep per-evaluation checkpoint files after a successful gather.",
+    )
+    evaluate.add_argument("--device-map", default="auto", help="Transformers device_map value.")
+    return parser
+
+
+def _read_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any]]]:
+    with path.open("r", encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            if line.strip():
+                value = json.loads(line)
+                if not isinstance(value, dict):
+                    raise ValueError(f"Line {line_number} must contain a JSON object")
+                yield line_number, value
+
+
+def _api_key(args: argparse.Namespace) -> str | None:
+    if not args.api_key_env:
+        return None
+    try:
+        return os.environ[args.api_key_env]
+    except KeyError as exc:
+        raise ValueError(f"Environment variable {args.api_key_env!r} is not set") from exc
+
+
+def _make_judge(args: argparse.Namespace):
+    key = _api_key(args)
+    if args.provider == "openai":
+        return OpenAIJudge(args.model, api_key=key)
+    if args.provider == "openai-compatible":
+        if not args.base_url:
+            raise ValueError("--base-url is required for --provider openai-compatible")
+        return OpenAICompatibleJudge(
+            args.model,
+            api_key=key,
+            base_url=args.base_url,
+            api_style=args.api_style,
+        )
+    if args.provider == "anthropic":
+        return AnthropicJudge(args.model, api_key=key)
+    return TransformersJudge(args.model, device_map=args.device_map)
+
+
+def _case_identifier(record: dict[str, Any], line_number: int) -> int | str:
+    for key in ("id", "index", "question"):
+        if key in record:
+            return record[key]
+    raise ValueError(f"Line {line_number} must contain one of: id, index, question")
+
+
+def _evaluate_record(
+    evaluator: GuidedBenchEvaluator,
+    line_number: int,
+    record: dict[str, Any],
+) -> tuple[dict[str, Any], float]:
+    if not isinstance(record.get("response"), str):
+        raise ValueError(f"Line {line_number} must contain a string response")
+    case = get_case(_case_identifier(record, line_number))
+    result = evaluator.evaluate(case, record["response"])
+    value = result.to_dict()
+    value["status"] = "completed"
+    value["input_line"] = line_number
+    for optional_key in ("method", "victim_model", "metadata"):
+        if optional_key in record:
+            value[optional_key] = record[optional_key]
+    return value, result.score
+
+
+@dataclass(frozen=True, slots=True)
+class _EvaluationTask:
+    task_id: str
+    line_number: int
+    record: dict[str, Any]
+    result_path: Path
+    error_path: Path
+
+
+def _checkpoint_directory(args: argparse.Namespace) -> Path:
+    if args.checkpoint_dir is not None:
+        return args.checkpoint_dir
+    return args.output.with_name(args.output.name + ".parts")
+
+
+def _task_id(
+    args: argparse.Namespace,
+    line_number: int,
+    record: dict[str, Any],
+) -> str:
+    payload = {
+        "package_version": __version__,
+        "benchmark_version": BENCHMARK_VERSION,
+        "provider": args.provider,
+        "model": args.model,
+        "base_url": args.base_url,
+        "api_style": args.api_style,
+        "max_tokens": args.max_tokens,
+        "input_line": line_number,
+        "record": record,
+    }
+    encoded = json.dumps(
+        payload,
+        ensure_ascii=False,
+        sort_keys=True,
+        separators=(",", ":"),
+    ).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()
+
+
+def _atomic_write_json(path: Path, value: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
+    try:
+        with temporary.open("x", encoding="utf-8") as handle:
+            handle.write(json.dumps(value, ensure_ascii=False) + "\n")
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary, path)
+    finally:
+        if temporary.exists():
+            temporary.unlink()
+
+
+def _atomic_write_jsonl(path: Path, values: Iterable[dict[str, Any]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
+    try:
+        with temporary.open("x", encoding="utf-8") as handle:
+            for value in values:
+                handle.write(json.dumps(value, ensure_ascii=False) + "\n")
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary, path)
+    finally:
+        if temporary.exists():
+            temporary.unlink()
+
+
+def _read_completed_result(path: Path, task_id: str) -> dict[str, Any] | None:
+    if not path.exists():
+        return None
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return None
+    if not isinstance(value, dict):
+        return None
+    if value.get("task_id") != task_id or value.get("status") != "completed":
+        return None
+    return value
+
+
+def _read_gathered_results(path: Path) -> dict[str, dict[str, Any]]:
+    if not path.exists():
+        return {}
+    results: dict[str, dict[str, Any]] = {}
+    for _, value in _read_jsonl(path):
+        task_id = value.get("task_id")
+        if isinstance(task_id, str) and value.get("status") == "completed":
+            results[task_id] = value
+    return results
+
+
+def _evaluate_and_checkpoint(
+    evaluator: GuidedBenchEvaluator,
+    task: _EvaluationTask,
+) -> tuple[bool, dict[str, Any]]:
+    try:
+        value, _ = _evaluate_record(evaluator, task.line_number, task.record)
+        value["task_id"] = task.task_id
+        _atomic_write_json(task.result_path, value)
+        if task.error_path.exists():
+            task.error_path.unlink()
+        return True, value
+    except Exception as exc:
+        failure = {
+            "task_id": task.task_id,
+            "status": "failed",
+            "input_line": task.line_number,
+            "identifier": _case_identifier(task.record, task.line_number),
+            "error_type": exc.__class__.__name__,
+            "error": str(exc),
+        }
+        _atomic_write_json(task.error_path, failure)
+        return False, failure
+
+
+def _clean_checkpoints(tasks: Iterable[_EvaluationTask], directory: Path) -> None:
+    for task in tasks:
+        for path in (task.result_path, task.error_path):
+            if path.exists():
+                path.unlink()
+    try:
+        directory.rmdir()
+    except OSError:
+        # Leave unrelated or crash-leftover files untouched.
+        pass
+
+
+def _run_evaluate(args: argparse.Namespace) -> int:
+    if args.workers < 1:
+        raise ValueError("--workers must be at least 1")
+    if args.provider == "transformers" and args.workers != 1:
+        raise ValueError(
+            "The local Transformers backend currently requires --workers 1; "
+            "serve the model behind an OpenAI-compatible endpoint for concurrent requests"
+        )
+
+    records = list(_read_jsonl(args.input))
+    for line_number, record in records:
+        if not isinstance(record.get("response"), str):
+            raise ValueError(f"Line {line_number} must contain a string response")
+        get_case(_case_identifier(record, line_number))
+    _api_key(args)  # Fail before scheduling work if an explicitly named key is missing.
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    checkpoint_dir = _checkpoint_directory(args)
+    tasks = tuple(
+        _EvaluationTask(
+            task_id=(task_id := _task_id(args, line_number, record)),
+            line_number=line_number,
+            record=record,
+            result_path=checkpoint_dir / f"{line_number:06d}-{task_id}.json",
+            error_path=checkpoint_dir / f"{line_number:06d}-{task_id}.error.json",
+        )
+        for line_number, record in records
+    )
+
+    gathered_cache = _read_gathered_results(args.output)
+    completed: dict[str, dict[str, Any]] = {}
+    pending: list[_EvaluationTask] = []
+    for task in tasks:
+        value = _read_completed_result(task.result_path, task.task_id)
+        if value is None:
+            value = gathered_cache.get(task.task_id)
+        if value is None:
+            pending.append(task)
+        else:
+            completed[task.task_id] = value
+
+    failures: list[dict[str, Any]] = []
+    if pending and args.workers == 1:
+        evaluator = GuidedBenchEvaluator(_make_judge(args), max_tokens=args.max_tokens)
+        for task in pending:
+            succeeded, value = _evaluate_and_checkpoint(evaluator, task)
+            if succeeded:
+                completed[task.task_id] = value
+            else:
+                failures.append(value)
+    elif pending:
+        worker_state = threading.local()
+
+        def evaluate_in_worker(
+            task: _EvaluationTask,
+        ) -> tuple[_EvaluationTask, bool, dict[str, Any]]:
+            evaluator = getattr(worker_state, "evaluator", None)
+            if evaluator is None:
+                evaluator = GuidedBenchEvaluator(
+                    _make_judge(args),
+                    max_tokens=args.max_tokens,
+                )
+                worker_state.evaluator = evaluator
+            succeeded, value = _evaluate_and_checkpoint(evaluator, task)
+            return task, succeeded, value
+
+        with ThreadPoolExecutor(max_workers=args.workers) as executor:
+            futures = [executor.submit(evaluate_in_worker, task) for task in pending]
+            for future in as_completed(futures):
+                task, succeeded, value = future.result()
+                if succeeded:
+                    completed[task.task_id] = value
+                else:
+                    failures.append(value)
+
+    if failures:
+        print(
+            json.dumps(
+                {
+                    "status": "incomplete",
+                    "completed": len(completed),
+                    "failed": len(failures),
+                    "checkpoint_dir": str(checkpoint_dir),
+                }
+            ),
+            file=sys.stderr,
+        )
+        return 1
+
+    missing = [task.task_id for task in tasks if task.task_id not in completed]
+    if missing:
+        raise RuntimeError(f"Missing {len(missing)} checkpoint results before gather")
+
+    ordered_results = [completed[task.task_id] for task in tasks]
+    _atomic_write_jsonl(args.output, ordered_results)
+    if not args.keep_checkpoints:
+        _clean_checkpoints(tasks, checkpoint_dir)
+
+    scores = [float(value["score"]) for value in ordered_results]
+    mean = sum(scores) / len(scores) if scores else 0.0
+    print(
+        json.dumps(
+            {
+                "status": "completed",
+                "evaluated": len(scores),
+                "recovered": len(tasks) - len(pending),
+                "mean_score": mean,
+                "workers": args.workers,
+                "output": str(args.output),
+            }
+        ),
+        file=sys.stderr,
+    )
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _parser().parse_args(argv)
+    try:
+        if args.command == "info":
+            all_cases = load_cases()
+            value = {
+                "package_version": __version__,
+                "benchmark_version": BENCHMARK_VERSION,
+                "cases": len(all_cases),
+                "core": len(load_cases(subset="core")),
+                "additional": len(load_cases(subset="additional")),
+                "topics": len({case.topic for case in all_cases}),
+            }
+            print(json.dumps(value, indent=2))
+            return 0
+        if args.command == "export":
+            path = export_jsonl(args.output, subset=args.subset, topic=args.topic)
+            print(path)
+            return 0
+        if args.command == "evaluate":
+            return _run_evaluate(args)
+    except (KeyError, MissingDependencyError, OSError, RuntimeError, ValueError) as exc:
+        print(f"guidedbench: error: {exc}", file=sys.stderr)
+        return 2
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

guidedbench/dataset.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import json
+from functools import lru_cache
+from pathlib import Path
+from typing import Literal
+
+from huggingface_hub import hf_hub_download
+from huggingface_hub.errors import GatedRepoError
+
+from ._version import BENCHMARK_VERSION
+from .models import GuidedBenchCase
+
+
+Subset = Literal["all", "core", "additional"]
+
+DATASET_REPO_ID = "HRXUST/GuidedBench"
+DATASET_REVISION = "10cc683c8df8c07139d3f900e381fb52738123c3"
+_DATA_FILES = (
+    ("core", "core.jsonl", 180),
+    ("additional", "additional.jsonl", 20),
+)
+
+
+@lru_cache(maxsize=8)
+def _all_cases(
+    repo_id: str,
+    revision: str,
+    cache_dir: str | None,
+    local_files_only: bool,
+) -> tuple[GuidedBenchCase, ...]:
+    cases: list[GuidedBenchCase] = []
+    try:
+        for expected_subset, filename, expected_count in _DATA_FILES:
+            data_file = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                repo_type="dataset",
+                revision=revision,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+            )
+            subset_cases: list[GuidedBenchCase] = []
+            with Path(data_file).open("r", encoding="utf-8") as handle:
+                for line_number, line in enumerate(handle, start=1):
+                    if not line.strip():
+                        continue
+                    try:
+                        case = GuidedBenchCase.from_dict(json.loads(line))
+                    except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
+                        raise RuntimeError(
+                            f"Invalid GuidedBench record in {filename}:{line_number}"
+                        ) from exc
+                    if case.subset != expected_subset:
+                        raise RuntimeError(
+                            f"Expected subset {expected_subset!r} in {filename}, "
+                            f"found {case.subset!r}"
+                        )
+                    subset_cases.append(case)
+            if len(subset_cases) != expected_count:
+                raise RuntimeError(
+                    f"Expected {expected_count} cases in {filename}, "
+                    f"found {len(subset_cases)}"
+                )
+            cases.extend(subset_cases)
+    except GatedRepoError as exc:
+        raise RuntimeError(
+            "GuidedBench dataset access is gated. Request access at "
+            f"https://huggingface.co/datasets/{repo_id}, then authenticate with "
+            "`hf auth login` or set HF_TOKEN."
+        ) from exc
+
+    if len(cases) != 200:
+        raise RuntimeError(f"Expected 200 GuidedBench cases, found {len(cases)}")
+    if len({case.id for case in cases}) != len(cases):
+        raise RuntimeError("GuidedBench contains duplicate case ids")
+    if any(case.benchmark_version != BENCHMARK_VERSION for case in cases):
+        raise RuntimeError(
+            f"Dataset revision {revision} does not match benchmark version "
+            f"{BENCHMARK_VERSION}"
+        )
+    return tuple(cases)
+
+
+def _load_source(
+    *,
+    repo_id: str,
+    revision: str,
+    cache_dir: str | Path | None,
+    local_files_only: bool,
+) -> tuple[GuidedBenchCase, ...]:
+    normalized_cache = (
+        str(Path(cache_dir).expanduser()) if cache_dir is not None else None
+    )
+    return _all_cases(repo_id, revision, normalized_cache, local_files_only)
+
+
+def load_cases(
+    *,
+    subset: Subset = "all",
+    topic: str | None = None,
+    repo_id: str = DATASET_REPO_ID,
+    revision: str = DATASET_REVISION,
+    cache_dir: str | Path | None = None,
+    local_files_only: bool = False,
+) -> tuple[GuidedBenchCase, ...]:
+    """Load a pinned GuidedBench release from the Hugging Face Hub cache."""
+
+    if subset not in {"all", "core", "additional"}:
+        raise ValueError("subset must be one of: all, core, additional")
+    cases = _load_source(
+        repo_id=repo_id,
+        revision=revision,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+    )
+    if subset != "all":
+        cases = tuple(case for case in cases if case.subset == subset)
+    if topic is not None:
+        cases = tuple(case for case in cases if case.topic == topic)
+    return cases
+
+
+def get_case(
+    identifier: int | str,
+    *,
+    repo_id: str = DATASET_REPO_ID,
+    revision: str = DATASET_REVISION,
+    cache_dir: str | Path | None = None,
+    local_files_only: bool = False,
+) -> GuidedBenchCase:
+    """Find a case by integer index, stable id, or exact question text."""
+
+    cases = _load_source(
+        repo_id=repo_id,
+        revision=revision,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+    )
+    for case in cases:
+        if isinstance(identifier, int) and case.index == identifier:
+            return case
+        if isinstance(identifier, str) and identifier in {case.id, case.question}:
+            return case
+    raise KeyError(f"No GuidedBench case matches {identifier!r}")
+
+
+def export_jsonl(
+    path: str | Path,
+    *,
+    subset: Subset = "all",
+    topic: str | None = None,
+    repo_id: str = DATASET_REPO_ID,
+    revision: str = DATASET_REVISION,
+    cache_dir: str | Path | None = None,
+    local_files_only: bool = False,
+) -> Path:
+    """Export records downloaded from Hugging Face as JSON Lines."""
+
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with destination.open("w", encoding="utf-8") as handle:
+        for case in load_cases(
+            subset=subset,
+            topic=topic,
+            repo_id=repo_id,
+            revision=revision,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        ):
+            handle.write(json.dumps(case.to_dict(), ensure_ascii=False) + "\n")
+    return destination

guidedbench/evaluator.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from .dataset import get_case
+from .judges.base import Judge
+from .models import EvaluationResult, GuidedBenchCase
+from .parsing import parse_guided_response
+from .prompt import JUDGE_SYSTEM_PROMPT, build_guided_prompt
+
+
+class GuidedBenchEvaluator:
+    """Evaluate responses using GuidedBench's verified case-specific guidelines."""
+
+    def __init__(self, judge: Judge, *, max_tokens: int = 800) -> None:
+        self.judge = judge
+        self.max_tokens = max_tokens
+
+    def evaluate(
+        self,
+        case: GuidedBenchCase | int | str,
+        response: str,
+    ) -> EvaluationResult:
+        resolved_case = case if isinstance(case, GuidedBenchCase) else get_case(case)
+        prompt = build_guided_prompt(resolved_case, response)
+        raw = self.judge.generate(
+            prompt,
+            system=JUDGE_SYSTEM_PROMPT,
+            max_tokens=self.max_tokens,
+        )
+        return parse_guided_response(raw, resolved_case, judge=self.judge.name)

guidedbench/judges/__init__.py

@@ -0,0 +1,13 @@
+from .anthropic import AnthropicJudge
+from .base import Judge, MissingDependencyError
+from .openai import OpenAICompatibleJudge, OpenAIJudge
+from .transformers import TransformersJudge
+
+__all__ = [
+    "AnthropicJudge",
+    "Judge",
+    "MissingDependencyError",
+    "OpenAICompatibleJudge",
+    "OpenAIJudge",
+    "TransformersJudge",
+]