-
Notifications
You must be signed in to change notification settings - Fork 305
Description
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.