Introduce Gemini Model Handler #35036
Conversation
|
Assigning reviewers: R: @liferoad for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
| """ A sample pipeline using the RunInference API to classify text using an LLM. | ||
| This pipeline creates a set of prompts and sends it to a Gemini service then | ||
| returns the predictions from the classifier model. This example uses the | ||
| gemini-2.0-flash-001 model. |
There was a problem hiding this comment.
probably need to update sdks/python/apache_beam/examples/inference/README.md with a new remote handler section.
| """ | ||
| if not isinstance(exception, errors.APIError): | ||
| return False | ||
| return exception.code == 429 or exception.code >= 500 |
There was a problem hiding this comment.
Shall we log these to make this more explicit to users?
There was a problem hiding this comment.
exceptions get logged by the RemoteModelHandler, no need to log them here (
)
| Args: | ||
| batch: a sequence of any values to be passed to the Gemini service. | ||
| Should be a list of strings. |
There was a problem hiding this comment.
Sequence[Any]: I assume we plan to support other types soon. So we use Any here, right?
There was a problem hiding this comment.
Ah yeah that was a missed line, I originally wrote this as string-only but decided it was better in the long-term to go ahead and make it more general. Let me update that docstring
2 participants
Introduces a basic Gemini Model Handler along with an example and tests. Currently provides a function for text prompts and responses, but is written to be extensible by allowing for custom request functions. The Gemini SDK contains a large amount of options, so this allows for the most general model handler possible (although post-processing to pull the responses out of the nested return type will still be necessary.)
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.