Skip to content

FEAT Add a WildGuard scorer following the LlamaGuard and ShieldGemma pattern - #2302

Open
immu4989 wants to merge 1 commit into
microsoft:mainfrom
immu4989:feat/wildguard-scorer
Open

FEAT Add a WildGuard scorer following the LlamaGuard and ShieldGemma pattern#2302
immu4989 wants to merge 1 commit into
microsoft:mainfrom
immu4989:feat/wildguard-scorer

Conversation

@immu4989

Copy link
Copy Markdown
Contributor

Closes #2265.

Third of the safety classifier set, after LlamaGuard (#1867) and ShieldGemma (#2261). WildGuard is a useful third because it judges a prompt and response together and returns three labels from one call, which is a different shape from the other two.

What it does

WildGuard answers three questions per call:

Harmful request: yes
Response refusal: no
Harmful response: yes

WildGuardLabel selects which one becomes the boolean score. The other two are kept in score_metadata, so reading them costs no extra request rather than three scorers repeating the same call.

Design notes

Prompt sourcing. The scored message is the response; the prompt it is judged against is read from the preceding turn of the scored conversation, or supplied with user_prompt=. It reads converted_value, since that is what the target actually received. This follows what we settled on in #2261.

N/A is a real value, not a parse failure. The paper documents N/A for the two response-side labels when no response was supplied, so the parser accepts and records it. If the selected label comes back N/A there is no boolean reading, so that does raise.

An empty response is rejected before the request, not in the parser. CallableResponseHandler converts any parser exception into InvalidJsonException, which drives a retry, and resending an empty response cannot change the N/A answer. So the scorer checks up front and points at WildGuardLabel.HARMFUL_REQUEST, which is tested to actually work with an empty response.

Template fidelity. The request reproduces the input format WildGuard was trained on (Table 12 of arXiv:2406.18495, which the paper states is also used at inference), asserted byte for byte in the tests. I omitted the chat scaffolding from AI2's reference implementation (<|user|>, [INST], [/INST], <|assistant|>), because PyRIT sends this as a user message and the serving layer applies its own template, so including it would wrap the request twice. Happy to change that if you would rather it match the reference string exactly.

On validation, which is the weak spot

I have not run this against a live WildGuard endpoint. The model is gated on HuggingFace, has no serverless inference provider, and is not in the Ollama library. So this is unit tested against the documented format rather than confirmed against the model. If you have a preferred hosting path I will run a live transcript and post it here.

Verification

pytest tests/unit   => 14773 passed, 5 skipped
pre-commit run      => all hooks pass, including ruff and ty

The prompt-sourcing, empty-response, and identity tests were each confirmed to fail without their fix.

Overlap with #2261

_resolve_user_prompt is close to the one in the ShieldGemma PR. I built this off main rather than stacking on #2261 so it is not blocked behind an unmerged branch. Once #2261 lands I am happy to factor the shared lookup into one helper in a follow-up.

The docs list also drops the hardcoded "Three"/"All three" count, since both PRs add an entry and it would otherwise need editing each time.

…pattern

Closes microsoft#2265.

WildGuard judges a user prompt and a model response together and returns three
labels from one call: whether the request is harmful, whether the response is a
refusal, and whether the response is harmful.

- `wildguard_parser.py` reads the three labelled lines. `N/A` is accepted, since
  the paper documents it as the value for the response-side labels when no
  response was supplied, rather than treating it as a malformed answer.
- `WildGuardLabel` selects which judgement becomes the boolean score. All three
  are kept in the score metadata, so reading the other two costs no extra
  request.
- `WildGuardScorer` scores a response and reads the prompt it is judged against
  from the preceding turn of the scored conversation, using the converted value
  the target actually received, or from a supplied `user_prompt`.
- An empty response is rejected before the request rather than in the parser,
  because parser exceptions drive a retry and resending cannot change the `N/A`
  answer.
- The request template reproduces the input format WildGuard was trained on
  (Table 12 of arXiv:2406.18495), asserted byte for byte in the tests. The chat
  scaffolding around it is omitted because the serving layer applies its own.
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.

FEAT Add a WildGuard scorer following the LlamaGuard and ShieldGemma pattern

1 participant