Skip to content

Add a dedicated ShieldGemma class #1973

@RyanMullins

Description

@RyanMullins

Is your feature request related to a problem? Please describe.

ShieldGemma was released in Jul 2024 with v0.14.3, but since then we've had feedback that the Example Use on the model card is prohibitively complex.

Describe the solution you'd like

I want to make a keras_hub.models.ShieldGemma class that encapsulates the functionality from the Example Use so that a user can load and use the model in two lines:

shieldgemma = keras_hub.models.ShieldGemma.from_preset("shieldgemma_2b_en")
shieldgemma.predict(inputs)

Describe alternatives you've considered

The primary alternative is to use ShieldGemma via the existing method from the Example Use.

Additional context

What remains unclear to me is precisely how to implement the keras_hub.models.ShieldGemma class. Should it be a subclass of keras.Model, keras_hub.models.Task, keras_hub.models.TextClassifier, or something else entirely?

I think the uncertainty here is that ShieldGemma is acting like a keras_hub.models.TextClassifier by using Gemma-as-a-Judge, but the architecture isn't quite what you could expect from a conventional classifier, like BERT. Instead of a dense layer over a limited vocabulary, the classification probability is computed from the logits for two specific tokens in the vocabulary.

Metadata

Metadata

Assignees

Labels

GemmaGemma model specific issues

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions