Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 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/trigger_files/beam_PostCommit_Python.json
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run.",
"modification": 29
"modification": 30
}

Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,15 @@
model can be generated by fine tuning the text-bison model or another similar
model (see
https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models#supervised-fine-tuning)

**NOTE**
Google has deprecated non-Gemini LLMs and no longer supports querying them on
Vertex AI endpoints. Separately, the use of the Vertex AI Predict API is not
supported for Gemini models in favor of use of the google-genai API. As a
result, this example no longer works. To perform inference with remote LLMs
deployed on Google infrastructure, please see the `GeminiModelHandler`
(in `apache_beam.ml.inference.gemini_inference`) and the
`gemini_text_classification.py` example.
"""

import argparse
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,15 +29,13 @@
# pylint: disable=ungrouped-imports
try:
from apache_beam.examples.inference import vertex_ai_image_classification
from apache_beam.examples.inference import vertex_ai_llm_text_classification
except ImportError as e:
raise unittest.SkipTest(
"Vertex AI model handler dependencies are not installed")

_INPUT = "gs://apache-beam-ml/testing/inputs/vertex_images/*/*.jpg"
_OUTPUT_DIR = "gs://apache-beam-ml/testing/outputs/vertex_images"
_FLOWER_ENDPOINT_ID = "5384055553544683520"
_LLM_ENDPOINT_ID = "9157860935048626176"
_ENDPOINT_PROJECT = "apache-beam-testing"
_ENDPOINT_REGION = "us-central1"
_ENDPOINT_NETWORK = "projects/844138762903/global/networks/beam-test-vpc"
Expand Down Expand Up @@ -65,21 +63,6 @@ def test_vertex_ai_run_flower_image_classification(self):
test_pipeline.get_full_options_as_args(**extra_opts))
self.assertEqual(FileSystems().exists(output_file), True)

@pytest.mark.vertex_ai_postcommit
def test_vertex_ai_run_llm_text_classification(self):
output_file = '/'.join([_OUTPUT_DIR, str(uuid.uuid4()), 'output.txt'])

test_pipeline = TestPipeline(is_integration_test=True)
extra_opts = {
'output': output_file,
'endpoint_id': _LLM_ENDPOINT_ID,
'endpoint_project': _ENDPOINT_PROJECT,
'endpoint_region': _ENDPOINT_REGION
}
vertex_ai_llm_text_classification.run(
test_pipeline.get_full_options_as_args(**extra_opts))
self.assertEqual(FileSystems().exists(output_file), True)


if __name__ == '__main__':
logging.getLogger().setLevel(logging.DEBUG)
Expand Down
Loading