FEAT Add a WildGuard scorer following the LlamaGuard and ShieldGemma pattern - #2302
Open
immu4989 wants to merge 1 commit into
Open
FEAT Add a WildGuard scorer following the LlamaGuard and ShieldGemma pattern#2302immu4989 wants to merge 1 commit into
immu4989 wants to merge 1 commit into
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
WildGuardLabelselects which one becomes the boolean score. The other two are kept inscore_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 readsconverted_value, since that is what the target actually received. This follows what we settled on in #2261.N/Ais a real value, not a parse failure. The paper documentsN/Afor the two response-side labels when no response was supplied, so the parser accepts and records it. If the selected label comes backN/Athere is no boolean reading, so that does raise.An empty response is rejected before the request, not in the parser.
CallableResponseHandlerconverts any parser exception intoInvalidJsonException, which drives a retry, and resending an empty response cannot change theN/Aanswer. So the scorer checks up front and points atWildGuardLabel.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
The prompt-sourcing, empty-response, and identity tests were each confirmed to fail without their fix.
Overlap with #2261
_resolve_user_promptis close to the one in the ShieldGemma PR. I built this offmainrather 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.