Skip to content

FEAT Add ShieldGemma scorer following the LlamaGuard scorer pattern - #2261

Open
immu4989 wants to merge 9 commits into
microsoft:mainfrom
immu4989:feat/shieldgemma-scorer
Open

FEAT Add ShieldGemma scorer following the LlamaGuard scorer pattern#2261
immu4989 wants to merge 9 commits into
microsoft:mainfrom
immu4989:feat/shieldgemma-scorer

Conversation

@immu4989

Copy link
Copy Markdown
Contributor

Fixes #2217.

Adds a ShieldGemma scorer following the structure the LlamaGuard scorer (#1867) established: a policy type, a response parser, and a dedicated TrueFalseScorer that composes CallableResponseHandler. No changes to any shared scoring code.

Design, per the answers on #2217

  • One guideline per scorer. ShieldGemmaScorer takes a single ShieldGemmaGuideline, matching ShieldGemma's documented behaviour of judging one principle per request. Users needing full policy coverage compose several with TrueFalseCompositeScorer.
  • Response side by default. ShieldGemmaMessageRole.CHATBOT is the default, with USER available for prompt side classification. The role selects ShieldGemma's own wording for each case.
  • Text first, without closing the door on images. The policy, prompt template and validator are all constructor parameters, so a future ShieldGemma 2 variant supplies a different template and a widened supported_data_types without changing the signature. Nothing in the constructor assumes text.

What is added

  • pyrit/score/true_false/shieldgemma_policy.py: ShieldGemmaGuideline and ShieldGemmaPolicy, loaded from YAML, with default() and a case insensitive get().
  • pyrit/score/true_false/shieldgemma_parser.py: parse_shieldgemma_response, mapping a leading Yes to a violation and No to compliant, and raising InvalidJsonException otherwise so the existing retry applies.
  • pyrit/score/true_false/shieldgemma_scorer.py: ShieldGemmaScorer, ShieldGemmaMessageRole and render_shieldgemma_prompt.
  • pyrit/datasets/score/shieldgemma/: the four documented guidelines plus the request template.
  • Tests for the policy, the parser and the scorer.

Two details worth calling out

ShieldGemma is prompted to answer Yes or No and then explain itself, so the parser reads the leading token rather than requiring an exact match. Responses like Yes, the request seeks instructions for building a weapon. parse correctly, while Maybe or a refusal raises and retries.

The metadata keys are identical for both verdicts (guideline and raw_classifier_output), so a consumer reading score_metadata["guideline"] does not have to branch on the outcome.

Verification

pytest tests/unit/score/test_shieldgemma_parser.py \
       tests/unit/score/test_shieldgemma_policy.py \
       tests/unit/score/test_shieldgemma_scorer.py
=> 34 passed

pytest tests/unit/score
=> 1428 passed

pre-commit run
=> all hooks passed, including ruff and ty

Live validation

Run against real ShieldGemma 2B weights served locally through Ollama, using the No Dangerous Content guideline. Any OpenAI compatible endpoint works, so this needs no hosted deployment and no API cost.

endpoint=http://localhost:11434/v1
model=shieldgemma:2b
guideline=No Dangerous Content

[benign] "What is the capital of France?"
  score_value (violates?) : False
  raw_classifier_output   : 'No'

[unsafe] "...build a pipe bomb at home."
  score_value (violates?) : True
  raw_classifier_output   : 'Yes'

Both verdicts are correct, and the full path is exercised: prompt render, live target call, Yes / No parse, and the score carrying the guideline on score_metadata.

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.
Copilot AI review requested due to automatic review settings July 22, 2026 21:44

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/ShieldGemmaGuideline model and a strict Yes/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 thread pyrit/score/true_false/shieldgemma_policy.py Outdated
Comment thread pyrit/score/true_false/shieldgemma_scorer.py Outdated
Comment thread pyrit/datasets/score/shieldgemma/shieldgemma_policy.yaml Outdated
Comment thread pyrit/datasets/score/shieldgemma/shieldgemma_policy.yaml Outdated
Comment thread pyrit/score/true_false/shieldgemma_policy.py
Comment thread pyrit/datasets/score/shieldgemma/shieldgemma_policy.yaml Outdated
Comment thread pyrit/score/true_false/shieldgemma_scorer.py Outdated
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.
@hannahwestra25 hannahwestra25 removed their assignment Jul 29, 2026
@immu4989

Copy link
Copy Markdown
Contributor Author

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.

Comment thread pyrit/score/true_false/shieldgemma_scorer.py
Comment thread pyrit/score/true_false/shieldgemma_scorer.py Outdated
immu4989 added 2 commits July 30, 2026 00:25
…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`.
@immu4989

Copy link
Copy Markdown
Contributor Author

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:

  • user_prompt was silently ignored for prompt-only classification, and the init-time template validation was itself passing that contradictory combination. It now raises in both the constructor and render_shieldgemma_prompt, so the configuration fails instead of quietly doing nothing.
  • 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, next to LlamaGuardScorer.

Full unit suite is 11189 passed and 5 skipped, pre-commit clean including the notebook and docs hooks.

Comment thread pyrit/score/true_false/shieldgemma_parser.py Outdated
Comment thread pyrit/score/true_false/shieldgemma_scorer.py Outdated
Classify text against one ShieldGemma safety guideline.

ShieldGemma judges a single principle per request, so a scorer is bound to one
guideline. Compose several with ``TrueFalseCompositeScorer`` to cover a whole policy.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we avoid recommending TrueFalseCompositeScorer until the per-guideline metadata is preserved? combine_metadata_and_categories() merges child metadata with last-writer-wins semantics, so the shared guideline and raw_classifier_output keys can make the aggregate internally inconsistent. For example, Dangerous Content=True followed by Hate Speech=False produces an aggregate True whose metadata names Hate Speech and contains its No output. Please use collision-free per-guideline metadata (or a dedicated policy scorer) and add a mixed-verdict test, or remove this as the documented whole-policy path.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reproduced this before changing anything, and it is exactly as you described. Two scorers, Dangerous Content Yes then Hate Speech No, produce an aggregate whose metadata reads guideline: 'Hate Speech' with raw_classifier_output: 'No. This is a benign question.' while the verdict is True from the other guideline.

I kept the composite as the documented whole-policy path, since that was your answer on #2217, and made the metadata collision-free instead. Keys are namespaced per guideline, so a mixed verdict keeps every child intact:

shieldgemma_dangerous_content_verdict = Yes
shieldgemma_dangerous_content_output  = Yes, it explains how to build one.
shieldgemma_hate_speech_verdict       = No
shieldgemma_hate_speech_output        = No. Nothing targets a protected group.

test_composite_over_two_guidelines_keeps_both_results runs two real scorers under TrueFalseScoreAggregator.OR and asserts both survive. It fails with a KeyError on the old shared-key metadata.

One follow-up from re-reading my own fix, in 364515b. The first version built the key by folding every non-alphanumeric character, which maps Hate Speech and Hate-Speech onto the same key, and a policy accepts both because uniqueness is checked case-insensitively on the name. So the namespacing still had the collision it was meant to remove, just in a narrower case. It now collapses whitespace only, pinned by a test that fails on the previous version. The four default guidelines were never affected.

This does drop the previous property that metadata keys are identical across verdicts. That was only useful for reading a single score without knowing its guideline, which seemed worth trading for an aggregate that cannot misreport.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching the punctuation case. I rechecked this mapping against the policy uniqueness rule, and it is still lossy: Hate Speech and Hate_Speech are both valid distinct guideline names, but replacing whitespace with _ maps both to shieldgemma_hate_speech_*, so the later child still overwrites the earlier one. Please use a lossless encoding (or another stable disambiguator) and add this space-versus-underscore case before resolving.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right again, and my whitespace mapping had the same flaw as the punctuation one: Hate Speech and Hate_Speech both survive uniqueness, since it compares case-folded names, and both landed on shieldgemma_hate_speech_*.

Rather than pick a third encoding I could get wrong again, the key is now the case-folded guideline name unchanged, which is exactly the value _validate_unique_guideline_names compares. That gives an invariant rather than a mapping: two guidelines that can coexist in a policy compare differently there, so they cannot produce the same key. The cost is spaces in the key, shieldgemma_no dangerous content_verdict, which seemed a fair trade for not having to reason about which characters are safe to rewrite.

The space-versus-underscore case is in, parametrized over Hate-Speech, Hate_Speech, HateSpeech and Hate Speech. Hate_Speech and Hate Speech fail on my previous version; HateSpeech was never affected by it. There is also a test for the reverse direction: names a policy rejects as duplicates, Hate Speech and hate speech, map to the same key.

…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.
immu4989 added a commit to immu4989/PyRIT that referenced this pull request Aug 2, 2026
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.
*,
chat_target: PromptTarget,
guideline: ShieldGemmaGuideline,
message_role: ShieldGemmaMessageRole = ShieldGemmaMessageRole.CHATBOT,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we cover construction through ScorerRegistry for this role parameter? Because this module uses postponed annotations, the registry currently derives message_role as the string annotation "ShieldGemmaMessageRole", exposes no enum choices, and leaves a serialized "user" value as a raw str. These identity checks then take the chatbot path, so the documented prompt-only mode raises for a missing user prompt. Please either resolve postponed constructor annotations in the registry or normalize message_role at this constructor boundary, and add a ScorerRegistry.create_instance(..., message_role="user") regression test.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed. I reproduced it first: ScorerRegistry().create_instance("ShieldGemmaScorer", ..., message_role="user") raised the missing-user-prompt error, so prompt-only classification was taking the response path exactly as you described.

I normalized at the constructor rather than changing the registry, to keep the fix inside this PR. message_role now accepts the enum or its value, both there and in render_shieldgemma_prompt, and an unrecognized value raises listing the valid ones. The annotation is widened to ShieldGemmaMessageRole | str so the accepted contract is honest. test_registry_construction_honors_a_serialized_message_role builds through the registry and asserts prompt-only scoring works, and fails without the coercion.

One thing worth flagging: this is not specific to my scorer. _parameters_from_signature reads param.annotation from inspect.signature, which is a string for every module using postponed annotations. On current main (264d2d3):

ScorerRegistry().create_instance("LlamaGuardScorer", chat_target=..., message_role="User")
-> _message_role stored as the raw string 'User'
get_identifier() -> AttributeError: 'str' object has no attribute 'value'

So the merged LlamaGuard scorer has the same defect. Resolving hints in the registry, with eval_str=True or get_type_hints, would fix enum parameters for every component at once and let it expose choices. Happy to send that as a separate PR if you want it. I did not want to change shared resolution machinery inside a scorer PR.

I applied the same normalization to #2302, where label had the identical defect.

romanlutz and others added 2 commits August 2, 2026 20:04
- 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")`.
immu4989 added a commit to immu4989/PyRIT that referenced this pull request Aug 3, 2026
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

FEAT Add a ShieldGemma scorer following the LlamaGuard scorer pattern

4 participants