FEAT Add ShieldGemma scorer following the LlamaGuard scorer pattern - #2261
FEAT Add ShieldGemma scorer following the LlamaGuard scorer pattern#2261immu4989 wants to merge 11 commits into
Conversation
Adds ShieldGemmaPolicy/ShieldGemmaGuideline, parse_shieldgemma_response, and a dedicated ShieldGemmaScorer composing CallableResponseHandler, plus the bundled policy and prompt YAMLs. One guideline per scorer per the design discussion on microsoft#2217; response-side classification by default; policy, template and validator are all injectable so image support can be added without a breaking change.
There was a problem hiding this comment.
Pull request overview
This PR introduces a new ShieldGemma-based true/false safety classifier scorer in pyrit/score/true_false, modeled after the existing LlamaGuard scorer architecture (policy + parser + dedicated scorer), and wires it into the public pyrit.score API with accompanying prompt/policy datasets and unit tests.
Changes:
- Adds
ShieldGemmaScorer(with message-role support) plus prompt rendering to classify a single message against a single guideline. - Adds a YAML-backed
ShieldGemmaPolicy/ShieldGemmaGuidelinemodel and a strictYes/No-leading-token response parser. - Adds default ShieldGemma prompt + policy assets and unit tests covering policy, parser, and scorer behavior.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/score/test_shieldgemma_scorer.py | Unit tests for prompt rendering, scoring behavior, metadata, and retry-on-parse-failure behavior. |
| tests/unit/score/test_shieldgemma_policy.py | Unit tests for YAML policy loading, rendering, case-insensitive lookup, and validation rules. |
| tests/unit/score/test_shieldgemma_parser.py | Unit tests for parsing Yes/No (including leading-token + reasoning cases) and malformed-response handling. |
| pyrit/score/true_false/shieldgemma_scorer.py | Implements the ShieldGemma true/false scorer and prompt rendering helpers. |
| pyrit/score/true_false/shieldgemma_policy.py | Implements guideline + policy models and YAML loading for default/custom ShieldGemma policies. |
| pyrit/score/true_false/shieldgemma_parser.py | Implements ShieldGemma response parsing into the CallableResponseHandler dictionary contract. |
| pyrit/score/init.py | Exposes ShieldGemma policy/parser/scorer symbols through the pyrit.score package API. |
| pyrit/datasets/score/shieldgemma/shieldgemma_prompt.yaml | Adds the default ShieldGemma request template as a seed prompt asset. |
| pyrit/datasets/score/shieldgemma/shieldgemma_policy.yaml | Adds the default ShieldGemma guideline set as a YAML policy asset. |
Response classification previously sent only the model response while the guidelines stayed worded for the prompt-only case, so requests matched neither of Google's documented modes. - Guidelines now store the role-independent body of a principle and take the attributing clause from the message role, so a guideline cannot be worded for one use case and sent in the other. - Response classification renders the Human Question alongside the Chatbot Response, reading the prompt from the preceding turn of the scored conversation or from a supplied user_prompt. - Added the prompt-response request template and restored the full Hate Speech guideline, both verbatim from Google's instruction. - Guideline name uniqueness is now checked case-insensitively, matching get(). - Guideline names are quoted unconditionally when rendered. - Prompt templates are validated when the scorer is constructed.
|
Thanks for the review. Pushed 3e239e8, which addresses everything here. Replies on your three threads for the details. The main change is that ShieldGemma now follows both of Google's documented use cases instead of a blend of the two. The underlying cause of the response mode problem was that guideline text and message role could drift apart, so a guideline now stores only the shared body of a principle and the role supplies the attributing clause. Across Google's two guideline tables that clause is the only thing that differs, so the mismatch cannot happen now. On the Copilot pass: guideline names are quoted unconditionally when rendered, and prompt templates are validated when the scorer is constructed rather than on the first scored message, both covered by tests. I did not take the two grammar suggestions on the policy YAML. "content that harming oneself" and "content that targeting identity" are verbatim from Google's model card, and since ShieldGemma was fine-tuned on those exact strings, correcting them moves the default policy off the distribution it was trained on. There is a comment at the top of the file recording that so it does not get tidied later. Both request templates are asserted byte for byte against the training instruction. Score suite is 1439 passed, pre-commit clean. |
…identity - `_resolve_user_prompt` reads `converted_value` from the preceding user turn instead of `original_value`. The converted value is what the target actually received, so after a converter runs the original can be the seed prompt and ShieldGemma would judge the response against context the target never saw. - A fixed `user_prompt` is now part of the component identity on the response side, so two scorers that send different requests no longer share an identifier and hash. Both are pinned by tests that fail without the change.
Self-review follow-ups found while re-reading the scorer against itself. - `user_prompt` was silently ignored for prompt-only classification. It now raises in both `render_shieldgemma_prompt` and the constructor, so a contradictory configuration fails instead of quietly doing nothing. - The init-time template validation was itself passing that contradictory combination, so it now only supplies a user prompt on the response side. - Corrected the `user_prompt` docstrings, which still described the old ignore-it behavior. - Added `ShieldGemmaScorer` to the external classifier integrations section of the true/false scorer docs, alongside `LlamaGuardScorer`.
|
Both fixed in abb5d4e, each pinned by a test that fails without the change (an AssertionError and a KeyError respectively). There is a follow-up in b83d00b from re-reading the scorer against itself rather than waiting for another round:
Full unit suite is 11189 passed and 5 skipped, pre-commit clean including the notebook and docs hooks. |
…ock the loop Addresses the third review round. - The parser dropped ShieldGemma's trailing explanation and replaced it with a generic sentence, contradicting its own contract. The explanation now flows into the rationale, which attacks consume as model feedback, and a verdict-only reply still falls back to the generic sentence. - Metadata keys are namespaced per guideline. `TrueFalseCompositeScorer` merges child metadata last-writer-wins, so the shared `guideline` and `raw_classifier_output` keys let a mixed-verdict aggregate report True while naming the guideline that returned False. A composite test now pins this. - Prompt resolution is async and reaches memory through `asyncio.to_thread`, so the blocking SQLAlchemy query no longer runs on the event loop during scoring. Each fix is covered by a test verified to fail without it.
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.
…e metadata keys Found while re-reading the namespacing. Folding every non-alphanumeric mapped "Hate Speech" and "Hate-Speech" onto one key, and a policy accepts both because uniqueness is checked case-insensitively on the name, so the collision this namespacing exists to prevent could still occur. Only whitespace is collapsed now, pinned by a test that fails on the previous slug.
- The metadata key now uses the case-folded guideline name unchanged, which is exactly what ShieldGemmaPolicy compares when it enforces unique names. Every rewriting was lossy: collapsing whitespace merged "Hate Speech" with "Hate_Speech", and folding punctuation also merged "Hate-Speech", yet a policy accepts each of those pairs. Deriving the key from the uniqueness value means two guidelines that can coexist can never share a key. - `message_role` is normalized at the constructor and in the public render helper. ScorerRegistry reads constructor annotations with `inspect.signature`, which under postponed annotations yields the string "ShieldGemmaMessageRole", so a serialized "user" arrived as a raw str, failed every identity check, and made prompt-only classification take the response path and raise. Both are pinned by tests that fail without the change, including `ScorerRegistry.create_instance(..., message_role="user")`.
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.
romanlutz
left a comment
There was a problem hiding this comment.
I think this is the last issue I'm seeing. Will take another thorough look after this is fixed.
| # TrueFalseCompositeScorer keeps every child's result. The aggregate merges child | ||
| # metadata last-writer-wins, so shared keys would leave it naming the wrong guideline. | ||
| "metadata": { | ||
| f"{slug}_verdict": leading_token.strip(_VERDICT_PUNCTUATION), |
There was a problem hiding this comment.
This still collides when a single Message contains multiple supported text pieces. TrueFalseScorer scores each piece and OR-aggregates the booleans, but every piece for this guideline emits these same two keys; the generic metadata merge is last-writer-wins. For example, ShieldGemma responses Yes then No return an aggregate True score whose retained *_verdict is No and whose *_output is only the second response. Could we preserve each piece's verdict/output under distinct keys (piece ID or index) and make any aggregate verdict reflect the configured aggregator? Please add a mixed-verdict multi-piece regression test as well.
… one A message can hold several supported text pieces. TrueFalseScorer scores each one and aggregates the booleans, but the metadata merge is last-writer-wins, so every piece emitted the same two keys and only the final piece survived. Pieces answering Yes then No produced an aggregate True whose retained verdict was No and whose output was only the second response. - Metadata keys now carry the scored piece's id, so each piece keeps its own verdict and raw output. - The guideline-level verdict is added after aggregation and follows the configured aggregator rather than whichever piece merged last. Pinned by a mixed-verdict multi-piece test that fails without either half.
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.
Follow-up to the multi-piece fix. Moving the memory read into a worker thread made it run concurrently, because pieces are scored under asyncio.gather, which put the shared SQLAlchemy session on several threads and made the lookup intermittently return nothing. A 40-trial probe failed 3 times. 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 carries no event-loop affinity. That removes the concurrency instead of serializing it and drops the redundant query per extra piece. 160 probe trials pass. A module-level asyncio.Lock was tried first and was wrong: it takes on the event loop that first uses it, so the same probe failed 39 times out of 40. The multi-piece test is also order independent now, since gather does not fix which piece is scored first.
Fixes #2217.
Adds a ShieldGemma scorer following the structure the LlamaGuard scorer (#1867) established: a policy type, a response parser, and a dedicated
TrueFalseScorerthat composesCallableResponseHandler. No changes to any shared scoring code.Design, per the answers on #2217
ShieldGemmaScorertakes a singleShieldGemmaGuideline, matching ShieldGemma's documented behaviour of judging one principle per request. Users needing full policy coverage compose several withTrueFalseCompositeScorer.ShieldGemmaMessageRole.CHATBOTis the default, withUSERavailable for prompt side classification. The role selects ShieldGemma's own wording for each case.supported_data_typeswithout changing the signature. Nothing in the constructor assumes text.What is added
pyrit/score/true_false/shieldgemma_policy.py:ShieldGemmaGuidelineandShieldGemmaPolicy, loaded from YAML, withdefault()and a case insensitiveget().pyrit/score/true_false/shieldgemma_parser.py:parse_shieldgemma_response, mapping a leadingYesto a violation andNoto compliant, and raisingInvalidJsonExceptionotherwise so the existing retry applies.pyrit/score/true_false/shieldgemma_scorer.py:ShieldGemmaScorer,ShieldGemmaMessageRoleandrender_shieldgemma_prompt.pyrit/datasets/score/shieldgemma/: the four documented guidelines plus the request template.Two details worth calling out
ShieldGemma is prompted to answer
YesorNoand then explain itself, so the parser reads the leading token rather than requiring an exact match. Responses likeYes, the request seeks instructions for building a weapon.parse correctly, whileMaybeor a refusal raises and retries.The metadata keys are identical for both verdicts (
guidelineandraw_classifier_output), so a consumer readingscore_metadata["guideline"]does not have to branch on the outcome.Verification
Live validation
Run against real ShieldGemma 2B weights served locally through Ollama, using the
No Dangerous Contentguideline. Any OpenAI compatible endpoint works, so this needs no hosted deployment and no API cost.Both verdicts are correct, and the full path is exercised: prompt render, live target call,
Yes/Noparse, and the score carrying the guideline onscore_metadata.