[yaml] - Add huggingface model handler by derrickaw · Pull Request #38110 · apache/beam

derrickaw · 2026-04-08T19:24:13Z

Please add a meaningful description for your change here

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: 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, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

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.

gemini-code-assist · 2026-04-08T19:24:21Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request extends the Apache Beam YAML support by integrating HuggingFace model handlers. It introduces a new provider class that allows users to leverage HuggingFace pipelines directly within their YAML-defined pipelines, facilitating easier machine learning model inference workflows.

Highlights

HuggingFace Integration: Added a new HuggingFacePipelineProvider to enable HuggingFace model handlers within the YAML pipeline framework.
Testing: Introduced a new YAML test case to verify the functionality of the HuggingFace pipeline integration.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

derrickaw · 2026-04-08T22:30:10Z

Blocked on #35198

derrickaw · 2026-04-15T11:54:11Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces support for HuggingFace pipelines in Beam YAML by implementing the HuggingFacePipelineProvider and adding a corresponding integration test. The review feedback suggests refining the HuggingFacePipelineProvider by using None as default values for task and model, updating the device type hint to support integer GPU indices, and implementing a validate method to ensure either a task or model is specified at configuration time.

derrickaw · 2026-05-05T01:06:54Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces support for HuggingFace pipelines within Beam YAML by implementing the HuggingFacePipelineProvider and adding corresponding integration tests. It also updates the HuggingFace inference module to handle TFAutoModel imports more robustly. A critical issue was identified in the HuggingFacePipelineProvider where task and model parameters were being defaulted to empty strings instead of None, which could cause runtime failures in the transformers library.

gemini-code-assist

Code Review

This pull request introduces support for HuggingFace pipelines in Beam YAML by implementing the HuggingFacePipelineProvider and adding a corresponding integration test. Review feedback identifies a missing Any import in the huggingface_inference.py module that could lead to a runtime error. Additionally, the custom inference_fn in the new YAML test requires corrections to properly handle batch processing and align with the expected output format for assertions.

gemini-code-assist

Code Review

This pull request introduces support for HuggingFace pipelines within Beam YAML by adding the HuggingFacePipelineProvider and an associated integration test. The review feedback identifies a parameter naming discrepancy in the provider class and a logic error in the test's inference function, which needs to return an iterable to correctly handle batch results.

github-actions · 2026-05-05T15:42:45Z

Assigning reviewers:

R: @damccorm for label python.
R: @damccorm for label build.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

damccorm · 2026-05-06T14:04:04Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces support for HuggingFace pipelines in the Beam YAML ML API by implementing a new HuggingFacePipelineProvider and adding robust import handling for TFAutoModel. The review identifies critical issues in the provided YAML test configuration, specifically regarding the implementation of the real_inference function and the subsequent data access pattern, providing necessary code corrections to ensure proper batch processing and result mapping.

damccorm

Thanks, had a couple minor comments, but generally LGTM

damccorm · 2026-05-06T14:05:52Z

        uses: ./.github/actions/gradle-command-self-hosted-action
        with:
-          gradle-command: :sdks:python:yamlIntegrationTests -PbeamPythonExtra=ml_test,yaml
+          gradle-command: :sdks:python:yamlIntegrationTests -PbeamPythonExtra=ml_test,yaml,transformers


Doesn't have to happen in this PR, but is there a reason we can't just install these as part of the task by default?

Relatedly, I think we need this here as well:

beam/.github/workflows/beam_PostCommit_Yaml_Xlang_Direct.yml

Line 83 in d6b98ba

gradle-command: :sdks:python:postCommitYamlIntegrationTests -PyamlTestSet=${{ matrix.test_set }} -PbeamPythonExtra=ml_test,yaml

damccorm · 2026-05-06T14:07:22Z

-from transformers import TFAutoModel
+
+try:
+  from transformers import TFAutoModel


What's special about TFAutoModel such that we now need to special case it, but not the other transformers imports?

codecov · 2026-05-06T18:21:49Z

Codecov Report

❌ Patch coverage is 30.00000% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 56.97%. Comparing base (efe4e94) to head (a033482).
⚠️ Report is 9 commits behind head on master.

Files with missing lines	Patch %	Lines
sdks/python/apache_beam/yaml/yaml_ml.py	30.00%	14 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #38110      +/-   ##
============================================
- Coverage     56.98%   56.97%   -0.01%     
  Complexity     3628     3628              
============================================
  Files          1185     1185              
  Lines        189433   189457      +24     
  Branches       3751     3751              
============================================
- Hits         107952   107948       -4     
- Misses        78024    78052      +28     
  Partials       3457     3457

Flag	Coverage Δ
python	`79.80% <30.00%> (-0.03%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

github-actions Bot added python yaml labels Apr 8, 2026

derrickaw force-pushed the 20260408_huggingFaceYaml branch 2 times, most recently from 17c61f2 to dd26577 Compare April 13, 2026 13:21

gemini-code-assist Bot reviewed Apr 15, 2026

View reviewed changes

Comment thread sdks/python/apache_beam/yaml/yaml_ml.py Outdated

Comment thread sdks/python/apache_beam/yaml/yaml_ml.py Outdated

derrickaw force-pushed the 20260408_huggingFaceYaml branch 2 times, most recently from f1431f5 to 0dad77b Compare April 15, 2026 13:34

github-actions Bot added the build label Apr 16, 2026

gemini-code-assist Bot reviewed May 5, 2026

View reviewed changes

Comment thread sdks/python/apache_beam/yaml/yaml_ml.py Outdated

derrickaw marked this pull request as ready for review May 5, 2026 14:21

gemini-code-assist Bot reviewed May 5, 2026

View reviewed changes

Comment thread sdks/python/apache_beam/ml/inference/huggingface_inference.py Outdated

Comment thread sdks/python/apache_beam/yaml/tests/runinference_huggingface.yaml

derrickaw marked this pull request as draft May 5, 2026 14:27

derrickaw marked this pull request as ready for review May 5, 2026 14:41

gemini-code-assist Bot reviewed May 5, 2026

View reviewed changes

Comment thread sdks/python/apache_beam/yaml/yaml_ml.py

Comment thread sdks/python/apache_beam/yaml/yaml_ml.py

Comment thread sdks/python/apache_beam/yaml/tests/runinference_huggingface.yaml

Comment thread sdks/python/apache_beam/yaml/tests/runinference_huggingface.yaml

github-actions Bot added the Next Action: Reviewers label May 5, 2026

gemini-code-assist Bot reviewed May 6, 2026

View reviewed changes

Comment thread sdks/python/apache_beam/yaml/tests/runinference_huggingface.yaml

Comment thread sdks/python/apache_beam/yaml/tests/runinference_huggingface.yaml

damccorm reviewed May 6, 2026

View reviewed changes

github-actions Bot added build and removed build labels May 6, 2026

derrickaw force-pushed the 20260408_huggingFaceYaml branch from 73b207a to 02ce797 Compare May 6, 2026 16:10

github-actions Bot added build and removed build labels May 6, 2026

derrickaw added 2 commits May 6, 2026 17:45

change runinference yaml name to vertexai

f792dd0

add huggingface support

129e3f9

derrickaw added 11 commits May 6, 2026 17:45

add huggingface test

6c62366

address gemini

e021846

fix lint issues

c9d1c2a

try adding transformers to testenv

52dccba

revert tox change and apply transformers to only the workflow needed

1132860

minor import and tweaks to get to work

c4dad0a

fix lint

d311bb8

updated test logic and model logic

c93432f

add new task with transformer dependency supplied

9e1e9e3

revert changes

bc1f8b5

try again without the try except

a033482

derrickaw force-pushed the 20260408_huggingFaceYaml branch from 02ce797 to a033482 Compare May 6, 2026 17:45

github-actions Bot added build and removed build labels May 6, 2026

Conversation

derrickaw commented Apr 8, 2026

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist Bot commented Apr 8, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

derrickaw commented Apr 8, 2026

Uh oh!

derrickaw commented Apr 15, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

derrickaw commented May 5, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented May 5, 2026

Uh oh!

damccorm commented May 6, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

damccorm left a comment

Choose a reason for hiding this comment

Uh oh!

damccorm May 6, 2026

Choose a reason for hiding this comment

Uh oh!

damccorm May 6, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented May 6, 2026

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants