File size: 1,497 Bytes
ca949ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """Dataset loading — bundled copy first, Hugging Face Hub as a fallback."""
from __future__ import annotations
import json
from functools import lru_cache
from pathlib import Path
from typing import List
from .types import AttackCase
DATA_DIR = Path(__file__).resolve().parents[1] / "data"
DATASET_REPO_ID = "NoahIsARider/RecRedTeam-ASI-Bench"
def _read_jsonl(path: Path) -> List[dict]:
rows = []
with path.open("r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
rows.append(json.loads(line))
return rows
def _resolve(filename: str) -> Path:
local = DATA_DIR / filename
if local.exists():
return local
try:
from huggingface_hub import hf_hub_download # noqa: PLC0415
return Path(hf_hub_download(DATASET_REPO_ID, f"data/{filename}", repo_type="dataset"))
except Exception as exc: # pragma: no cover - network / auth failures
raise FileNotFoundError(
f"Could not locate {filename} locally or on the Hub ({DATASET_REPO_ID}): {exc}"
) from exc
@lru_cache(maxsize=1)
def load_cases(split: str = "public") -> List[AttackCase]:
rows = _read_jsonl(_resolve("attacks.jsonl"))
cases = [AttackCase.from_dict(r) for r in rows]
return [c for c in cases if c.split == split]
@lru_cache(maxsize=1)
def load_taxonomy() -> List[dict]:
return _read_jsonl(_resolve("taxonomy.jsonl"))
|