data_utils.py

"""
Data loading utilities for AB-GCG experiments.
"""

import csv
import random
from pathlib import Path
from typing import List, Tuple


def load_advbench(
    path: str = "data/harmful_behaviors.csv",
    n: int = 50,
    max_words: int = 30,
    seed: int = 42,
) -> List[Tuple[str, str]]:
    """
    Load AdvBench harmful behaviors dataset.

    Args:
        path: Path to harmful_behaviors.csv.
        n: Number of prompts to sample.
        max_words: Maximum word count per prompt (filter long ones).
        seed: Random seed for sampling.

    Returns:
        List of (prompt, target) tuples.
    """
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(
            f"AdvBench file not found at {path}. "
            f"Download from: https://github.com/llm-attacks/llm-attacks/blob/main/data/advbench/harmful_behaviors.csv"
        )

    behaviors = []
    with open(p, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            goal = row.get("goal", "").strip()
            target_str = row.get("target", "").strip()
            if goal and len(goal.split()) <= max_words:
                behaviors.append((goal, target_str))

    if len(behaviors) < n:
        print(
            f"Warning: only {len(behaviors)} prompts with <= {max_words} words. "
            f"Using all of them instead of {n}."
        )
        n = len(behaviors)

    random.seed(seed)
    sampled = random.sample(behaviors, n)

    return sampled