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 6 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 |
Carries over the multi-piece review point from microsoft#2261 and fixes a race the earlier asyncio.to_thread change exposed. - Metadata keys now carry the scored piece's id, so each piece keeps its own labels and raw output instead of the last one overwriting the rest, and the label-level verdict added after aggregation follows the configured aggregator. - The memory lookup is serialized behind a module-level lock. Pieces are scored with asyncio.gather and TrueFalseCompositeScorer gathers its children, so moving the read into a worker thread put the shared SQLAlchemy session on several threads at once and the lookup intermittently returned nothing. A 40-trial probe failed 3 times without the lock and 160 times out of 160 with it. Pinned by a mixed-verdict multi-piece test, written to be order independent because gather does not fix which piece is scored first.
…r piece Replaces the module-level lock from the previous commit, which was wrong. An asyncio.Lock created at import time takes on the event loop that first uses it, so it fails once scoring runs on a different loop. A probe that gives each trial its own loop failed 39 times out of 40 with it. The user prompt belongs to the conversation rather than to an individual piece, so it is now resolved once per scored message and handed to the pieces through a ContextVar, which gather copies into each child task and which has no event-loop affinity. That removes the concurrent memory reads instead of serializing them, and drops the redundant query per extra piece. 160 probe trials pass with this and the earlier per-piece read failed 3 in 40.
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.