Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a dedicated ShieldGemma class #1973

Open
RyanMullins opened this issue Nov 5, 2024 · 1 comment
Open

Add a dedicated ShieldGemma class #1973

RyanMullins opened this issue Nov 5, 2024 · 1 comment
Assignees
Labels
Gemma Gemma model specific issues

Comments

@RyanMullins
Copy link
Contributor

RyanMullins commented Nov 5, 2024

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.

@github-actions github-actions bot added the Gemma Gemma model specific issues label Nov 5, 2024
@RyanMullins
Copy link
Contributor Author

Tagging @mattdangerw who helped with the ShieldGemma launch.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Gemma Gemma model specific issues
Projects
None yet
Development

No branches or pull requests

4 participants