FEAT Add a WildGuard scorer following the LlamaGuard and ShieldGemma pattern - #2302
FEAT Add a WildGuard scorer following the LlamaGuard and ShieldGemma pattern#2302immu4989 wants to merge 4 commits 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.
Carries over a review point from microsoft#2261, which shares this lookup. Prompt resolution is now async and reaches memory through `asyncio.to_thread`, so the blocking SQLAlchemy query no longer runs on the event loop during scoring.
…d path One scorer already reports all three judgements, so a second only repeats the same request and the two scores would carry the same metadata keys.
|
Carried a review point over from #2261 in b586d54: prompt resolution is async and reaches memory through One related question while it is in your head. You flagged on #2261 that I have not namespaced these keys, because composing is not the intended path here. A single WildGuard call returns all three judgements, so a second scorer only repeats the same request to read a field the first one already reported. I documented that in the class docstring rather than changing the schema. Happy to namespace them the same way as ShieldGemma if you would rather the rule be uniform across the classifier scorers. Your call, I did not want to guess and churn the schema. |
Carries over a review point from microsoft#2261, which has the same defect. ScorerRegistry reads constructor annotations with `inspect.signature`, which under postponed annotations yields the string "WildGuardLabel", so a configured `label="Harmful request"` arrived as a raw str. That failed the identity guard for the empty-response check and missed the parser's per-label lookup, so scoring raised instead of selecting the requested judgement. Pinned by a `ScorerRegistry.create_instance(..., label="Harmful request")` test plus an unknown-label case.
|
Carried the Built through the registry, Pinned by a |
2 participants
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.