Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/beam_PostCommit_Yaml_Xlang_Direct.yml
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ jobs:
- name: run PostCommit Yaml Xlang Direct script
uses: ./.github/actions/gradle-command-self-hosted-action
with:
gradle-command: :sdks:python:postCommitYamlIntegrationTests -PyamlTestSet=${{ matrix.test_set }}
gradle-command: :sdks:python:postCommitYamlIntegrationTests -PyamlTestSet=${{ matrix.test_set }} -PbeamPythonExtra=ml_test,yaml
- name: Archive Python Test Results
uses: actions/upload-artifact@v7
if: failure()
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/beam_PreCommit_Yaml_Xlang_Direct.yml
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ jobs:
- name: run PreCommit Yaml Xlang Direct script
uses: ./.github/actions/gradle-command-self-hosted-action
with:
gradle-command: :sdks:python:yamlIntegrationTests
gradle-command: :sdks:python:yamlIntegrationTests -PbeamPythonExtra=ml_test,yaml
- name: Archive Python Test Results
uses: actions/upload-artifact@v7
if: failure()
Expand Down
62 changes: 0 additions & 62 deletions sdks/python/apache_beam/yaml/tests/runinference_huggingface.yaml

This file was deleted.

49 changes: 0 additions & 49 deletions sdks/python/apache_beam/yaml/yaml_ml.py
Original file line number Diff line number Diff line change
Expand Up @@ -282,55 +282,6 @@ def inference_output_type(self):
('model_id', Optional[str])])


@ModelHandlerProvider.register_handler_type('HuggingFacePipeline')
class HuggingFacePipelineProvider(ModelHandlerProvider):
def __init__(
self,
task: Optional[str] = None,
model: Optional[str] = None,
preprocess: Optional[dict[str, str]] = None,
postprocess: Optional[dict[str, str]] = None,
device: Optional[Any] = None,
inference_fn: Optional[dict[str, str]] = None,
load_pipeline_args: Optional[dict[str, Any]] = None,
**kwargs):
try:
from apache_beam.ml.inference.huggingface_inference import HuggingFacePipelineModelHandler
except ImportError:
raise ValueError(
'Unable to import HuggingFacePipelineModelHandler. Please '
'install transformers dependencies.')

kwargs = {k: v for k, v in kwargs.items() if not k.startswith('_')}

inference_fn_obj = self.parse_processing_transform(
inference_fn, 'inference_fn') if inference_fn else None

handler_kwargs = {}
if inference_fn_obj:
handler_kwargs['inference_fn'] = inference_fn_obj

_handler = HuggingFacePipelineModelHandler(
task=task,
model=model,
device=device,
load_pipeline_args=load_pipeline_args,
**handler_kwargs,
**kwargs)

super().__init__(_handler, preprocess, postprocess)

@staticmethod
def validate(config):
if not config.get('task') and not config.get('model'):
raise ValueError(
"HuggingFacePipeline requires either 'task' or "
"'model' to be specified.")

def inference_output_type(self):
return Any


@beam.ptransform.ptransform_fn
def run_inference(
pcoll,
Expand Down
14 changes: 2 additions & 12 deletions sdks/python/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -124,20 +124,10 @@ tasks.register("generateYamlDocs") {
outputs.file "${buildDir}/yaml-examples.html"
}

tasks.register("installYamlIntegrationTestDeps") {
dependsOn installGcpTest
doLast {
exec {
executable 'sh'
args '-c', ". ${envdir}/bin/activate && pip install --pre --retries 10 ${buildDir}/apache-beam.tar.gz[ml_test,yaml,transformers]"
}
}
}

tasks.register("yamlIntegrationTests") {
description "Runs precommit integration tests for yaml pipelines."

dependsOn installYamlIntegrationTestDeps
dependsOn installGcpTest
// Need to build all expansion services referenced in apache_beam/yaml/*.*
// grep -oh 'sdk.*Jar' sdks/python/apache_beam/yaml/*.yaml | sort | uniq
dependsOn ":sdks:java:extensions:schemaio-expansion-service:shadowJar"
Expand All @@ -156,7 +146,7 @@ tasks.register("yamlIntegrationTests") {
tasks.register("postCommitYamlIntegrationTests") {
description "Runs postcommit integration tests for yaml pipelines - parameterized by yamlTestSet."

dependsOn installYamlIntegrationTestDeps
dependsOn installGcpTest
// Need to build all expansion services referenced in apache_beam/yaml/*.*
// grep -oh 'sdk.*Jar' sdks/python/apache_beam/yaml/*.yaml | sort | uniq
dependsOn ":sdks:java:extensions:schemaio-expansion-service:shadowJar"
Expand Down
Loading