Skip to content

FEAT: Add remote WildGuardMix dataset loader#1827

Merged
romanlutz merged 12 commits into
microsoft:mainfrom
romanlutz:romanlutz/jailbreak-dataset-loaders
Jul 18, 2026
Merged

FEAT: Add remote WildGuardMix dataset loader#1827
romanlutz merged 12 commits into
microsoft:mainfrom
romanlutz:romanlutz/jailbreak-dataset-loaders

Conversation

@romanlutz

@romanlutz romanlutz commented May 28, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds a remote seed-dataset loader for AllenAI's WildGuardMix moderation corpus. The previously proposed TrustAIRLab remote loader has been removed from this PR; that corpus is being evaluated separately as a static, credited jailbreak-template collection.

The AdvBench investigation also confirmed that no new remote AdvBench loader is needed. Its separate harmful_strings.csv gap is tracked in #1826.

WildGuardMix loader

Loader Source Rows Size License Access
_WildGuardMixDataset allenai/wildguardmix 86,759 train + 1,725 test huge ODC-BY 1.0 Hugging Face token plus acceptance of the separate AI2 Responsible Use Guidelines

The loader:

  • Fetches both wildguardtrain and wildguardtest configs by default, with a splits filter to narrow the result.
  • Defaults to harmful, adversarial jailbreak prompts. Vanilla prompts, benign companions, and train rows containing responses remain explicit opt-ins through prompt_harm_labels, adversarial, and prompt_only.
  • Maps all 14 harmful upstream subcategory slugs to canonical PyRIT HarmCategory names while preserving the raw source label in seed metadata. Explicitly unharmful rows remain uncategorized.
  • Is tagged {safety, jailbreak} and intentionally not tagged default.
  • Accepts an explicit token or falls back to HUGGINGFACE_TOKEN.

AdvBench investigation

PyRIT's local adv_bench.prompt contains 520 rows with harm-category metadata and matches the canonical llm-attacks/data/advbench/harmful_behaviors.csv goals. The viable Hugging Face mirror contains only 500 rows, so replacing or duplicating the local dataset with a remote loader would regress coverage.

Canonical AdvBench also includes 574 output-side strings in harmful_strings.csv, which serve a different scorer-evaluation use case. That gap is tracked separately in #1826.

Validation

  • 36 focused WildGuardMix unit tests cover both configs, filtering, empty results, enum validation, metadata, all upstream harm-category mappings, and token forwarding.
  • All 1,181 dataset unit tests pass after removing TrustAIRLab from this branch.
  • Daily dataset e2e discovery includes WildGuardMix and skips cleanly when no token is configured or the account has not accepted the dataset terms.
  • The complete set of 14 harmful source subcategories was checked against a public WildGuardMix mirror; unit-test prompts are fabricated mocks rather than copied gated rows.

Documentation and licensing

  • Adds the han2024wildguard citation and includes WildGuard in the built-in dataset documentation and rendered provider list.
  • WildGuardMix data is licensed under ODC-BY 1.0. The AI2 Responsible Use Guidelines are a separate Hugging Face access gate.
  • PyRIT does not redistribute the corpus; the loader fetches it at runtime after the user obtains access.

Follow-ups

romanlutz and others added 9 commits May 27, 2026 22:34
- _WildGuardMixDataset: AI2 WildGuardMix (gated, requires
  HUGGINGFACE_TOKEN). Single loader returns both wildguardtrain +
  wildguardtest by default. Opt-in filters for splits, prompt-harm
  label, adversarial flag, and prompt-only rows.
- _TrustAIRLabInTheWildJailbreakDataset: 1,403 in-the-wild
  jailbreak prompts (Shen et al. 2024, MIT licensed). Uses the
  jailbreak_2023_12_25 superset config. Opt-in platform filter
  (discord/reddit/website) and dedup toggle.
- Both tagged {safety, jailbreak} (jailbreak is a new tag).
- Daily e2e skips WildGuardMix when HUGGINGFACE_TOKEN is unset.
- 32 new unit tests (mocked _fetch_from_huggingface). All pass.
- Live HF sanity check on TrustAIRLab succeeded (1,403 seeds
  across {discord, reddit, website}); WildGuardMix auth path
  validated (token forwarded, gated 403 due to RUG-acceptance
  needed on the HF account).
- doc/references.bib: added han2024wildguard.
- doc/bibliography.md + loading_datasets.py: cite both papers.
- AdvBench investigated: PyRIT local is a strict superset of
  walledai/AdvBench (520 vs 500, plus Claude-labeled harm cats).
  No remote loader needed for behaviors. The 574-row
  harmful_strings.csv gap is filed as a follow-up issue (microsoft#1826).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Ran e2e parametrized test for both new loaders.
  TrustAIRLab passed (1,403 seeds live). WildGuardMix skipped
  because my HF account hasn't accepted AI2's RUG click-through.
- Verified the WildGuardMix row schema my loader assumes
  (prompt, response, prompt_harm_label, response_harm_label,
  response_refusal_label, adversarial, subcategory) against
  AllenAI's own production safety-eval code (allenai/safety-eval
  evaluation/tasks/{generation,classification}/wildguardtest/).
  All seven columns match upstream 1:1.
- Corrected WildGuardMix license docstring: ODC-BY 1.0 is the
  data license; AI2 Responsible Use Guidelines is the access
  gate, not the license.
- Hardened e2e skip: in addition to the HUGGINGFACE_TOKEN-set
  pre-check, catch 'gated dataset' errors from tokens that
  haven't accepted the RUG so CI doesn't fail on per-account
  gating instead of a real loader bug.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The vanilla (non-adversarial) WildGuardMix harmful prompts overlap heavily with what AdvBench already provides as plain `do X'' objectives. The point of shipping WildGuardMix as a jailbreak-tagged loader is the adversarial/jailbreak-style subset (DAN, persona wraps, etc.), so default `adversarial` to `[ADVERSARIAL]` instead of both values. Users who want the vanilla subset can pass `adversarial=[VANILLA]` or both explicitly.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ataset-loaders

# Conflicts:
#	doc/bibliography.md
#	pyrit/datasets/seed_datasets/remote/__init__.py
…ataset-loaders

# Conflicts:
#	doc/bibliography.md
#	doc/code/datasets/1_loading_datasets.ipynb
…ataset-loaders

# Conflicts:
#	doc/bibliography.md
#	doc/code/datasets/1_loading_datasets.ipynb
#	tests/end_to_end/test_all_datasets.py
…ataset-loaders

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: bd8f4525-ab6e-474a-9a3e-4d01c9f5b691

# Conflicts:
#	doc/bibliography.md
#	doc/code/datasets/1_loading_datasets.ipynb
#	pyrit/datasets/seed_datasets/remote/__init__.py
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: bd8f4525-ab6e-474a-9a3e-4d01c9f5b691
@romanlutz romanlutz changed the title FEAT: Add remote loaders for WildGuardMix and TrustAIRLab in-the-wild jailbreaks FEAT: Add remote WildGuardMix dataset loader Jul 16, 2026
…ataset-loaders

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: bd8f4525-ab6e-474a-9a3e-4d01c9f5b691

# Conflicts:
#	doc/code/datasets/1_loading_datasets.ipynb
@adrian-gavrila adrian-gavrila self-assigned this Jul 17, 2026
…ataset-loaders

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: bd8f4525-ab6e-474a-9a3e-4d01c9f5b691

# Conflicts:
#	doc/code/datasets/1_loading_datasets.ipynb

@adrian-gavrila adrian-gavrila left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved with one comment on tests, feel free to resolve if you disagree there and merge

Comment thread tests/unit/datasets/test_wildguardmix_dataset.py
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@romanlutz
romanlutz added this pull request to the merge queue Jul 18, 2026
Merged via the queue into microsoft:main with commit 8a91d7d Jul 18, 2026
53 checks passed
@romanlutz
romanlutz deleted the romanlutz/jailbreak-dataset-loaders branch July 18, 2026 04:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants