Skip to content

FEAT Add a ShieldGemma scorer following the LlamaGuard scorer pattern #2217

Description

@immu4989

Now that the LlamaGuard scorer (#1867) has landed as a dedicated TrueFalseScorer composing CallableResponseHandler, the same shape extends naturally to Google's ShieldGemma, the other widely used open-weight safety classifier. Opening this to propose the design and settle a couple of questions before writing code.

One important difference from LlamaGuard

ShieldGemma does not use LlamaGuard's safe / unsafe\n<codes> format. It is prompted with the content plus a single safety guideline and answers with Yes or No to "Does the content violate the above principle?". So the classifier decides one guideline per call rather than returning a multi-category verdict in one shot. The canonical text guidelines are a small set (dangerous content, harassment, hate speech, sexually explicit information), each a named principle with a description.

That makes the parser simpler than LlamaGuard's, but the policy model and prompt are different, so this is an adaptation of the pattern rather than a copy.

Proposed shape, mirroring the LlamaGuard files

  • pyrit/score/true_false/shieldgemma_policy.py: a ShieldGemmaGuideline (name plus description) and a ShieldGemmaPolicy (a versioned set of guidelines), parallel to LlamaGuardPolicy.
  • pyrit/score/true_false/shieldgemma_parser.py: parse_shieldgemma_response mapping a leading Yes to score_value="True" (violates) and No to "False", raising InvalidJsonException on anything else so the retry path applies.
  • pyrit/score/true_false/shieldgemma_scorer.py: ShieldGemmaScorer(TrueFalseScorer) composing CallableResponseHandler(parse_shieldgemma_response) and going through _run_llm_scoring_async, exactly like LlamaGuardScorer.
  • pyrit/datasets/score/shieldgemma/: a default policy YAML with the canonical guidelines plus the prompt template(s).
  • Tests mirroring test_llamaguard_parser, test_llamaguard_policy, and test_llamaguard_scorer.

Open design questions

  1. Guideline granularity. ShieldGemma is documented as most accurate with one guideline per call, which means N calls to cover N guidelines. The alternatives are a single scorer bound to one guideline (compose several for full coverage), or one scorer that runs the guidelines and aggregates. LlamaGuard did it all in one call; ShieldGemma cannot. Which default do you prefer?
  2. Prompt-versus-response classification. ShieldGemma has distinct templates for classifying a user prompt versus a model response, similar in spirit to LlamaGuardMessageRole. Expose that as a mode on the scorer?
  3. Version and modality. Bundle the Gemma-2 text policy first, and leave ShieldGemma-2 (image) as a follow-on, consistent with how the LlamaGuard PR scoped multimodal out?

I would validate the final scorer end to end against a hosted ShieldGemma endpoint the same way I did for LlamaGuard, and I can take ShieldGemma and then WildGuard as the natural continuation of the safety-classifier set. Happy to start once the guideline-granularity question in particular is settled.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions