FEAT Add ShieldGemma scorer following the LlamaGuard scorer pattern#2261
Open
immu4989 wants to merge 1 commit into
Open
FEAT Add ShieldGemma scorer following the LlamaGuard scorer pattern#2261immu4989 wants to merge 1 commit into
immu4989 wants to merge 1 commit 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.
Contributor
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. |
Comment on lines
+46
to
+47
| name = self.name if self.name.endswith((".", ":", "!", "?")) else f'"{self.name}"' | ||
| return f"{name}: {self.description}" |
Comment on lines
+135
to
+138
| self._prompt_template = _resolve_prompt_template(prompt_template=prompt_template) | ||
| self._response_handler = CallableResponseHandler( | ||
| parser=partial(parse_shieldgemma_response, guideline_name=guideline.name) | ||
| ) |
| guidelines: | ||
| - name: No Dangerous Content | ||
| description: >- | ||
| The prompt shall not contain or seek generation of content that harming oneself |
| physical threats, denial of tragic events, disparaging victims of violence). | ||
| - name: No Hate Speech | ||
| description: >- | ||
| The prompt shall not contain or seek generation of content that targeting identity |
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.
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.