diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json b/.github/trigger_files/beam_PostCommit_Python_Dependency.json index f3bc9e958674..4ff9a5d35f33 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Dependency.json +++ b/.github/trigger_files/beam_PostCommit_Python_Dependency.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 9 + "modification": 13 } diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py index 261b480b1083..b631b2b263bb 100644 --- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py +++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py @@ -72,7 +72,8 @@ def test_mltransform_compute_and_apply_vocab(self, mock_stdout): mltransform_compute_and_apply_vocabulary() predicted = mock_stdout.getvalue().splitlines() expected = check_mltransform_compute_and_apply_vocab() - self.assertEqual(predicted, expected) + # Print order from the pipeline is not guaranteed. + self.assertCountEqual(predicted, expected) def test_mltransform_scale_to_0_1(self, mock_stdout): mltransform_scale_to_0_1() @@ -84,7 +85,8 @@ def test_mltransform_compute_and_apply_vocab_scalar(self, mock_stdout): mltransform_compute_and_apply_vocabulary_with_scalar() predicted = mock_stdout.getvalue().splitlines() expected = check_mltransform_compute_and_apply_vocabulary_with_scalar() - self.assertEqual(predicted, expected) + # Print order from the pipeline is not guaranteed. + self.assertCountEqual(predicted, expected) if __name__ == '__main__': diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py index 56a9535b5764..f5e5342de53a 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py @@ -215,7 +215,14 @@ def run_inference( self._config.api_url, headers=self._config.authorization_token, json={"inputs": batch}) - return response.json() + response.raise_for_status() + result = response.json() + # Reject {"error": ...} so it is not treated as embedding values. + if isinstance(result, dict) and 'error' in result: + raise RuntimeError( + 'Hugging Face Inference API error for %s: %s' % + (self._config.api_url, result['error'])) + return result class InferenceAPIEmbeddings(EmbeddingsManager): diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py index 166827f3e907..ae8bd4ff8f14 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py @@ -27,6 +27,7 @@ import jsonpickle import numpy as np import pytest +import requests from parameterized import parameterized import apache_beam as beam @@ -70,7 +71,33 @@ # InferenceAPIEmbeddings calls /pipeline/feature-extraction. INFERENCE_API_MODEL_NAME = "BAAI/bge-small-en-v1.5" INFERENCE_API_EMBEDDING_DIM = 384 +INFERENCE_API_URL = ( + 'https://router.huggingface.co/hf-inference/models/' + f'{INFERENCE_API_MODEL_NAME}/pipeline/feature-extraction') IMAGE_MODEL_NAME = "clip-ViT-B-32" + + +def _hf_inference_api_available(): + """Returns (ready, reason) for the optional live Inference API dependency.""" + if not _HF_TOKEN: + return False, 'HF_INFERENCE_TOKEN environment variable not set.' + try: + response = requests.post( + INFERENCE_API_URL, + headers={'Authorization': f'Bearer {_HF_TOKEN}'}, + json={'inputs': test_query}, + timeout=60) + payload = response.json() + except Exception as exc: # pylint: disable=broad-except + return False, 'HF Inference API unavailable: %s' % exc + if response.status_code >= 400: + return False, 'HF Inference API HTTP %s: %s' % ( + response.status_code, payload) + if isinstance(payload, dict) and 'error' in payload: + return False, 'HF Inference API error: %s' % payload['error'] + return True, '' + + _parameterized_inputs = [ ([{ test_query_column: 'That is a happy person' @@ -476,7 +503,8 @@ def test_loading_model_from_artifacts_written_by_older_versions(self): self.assertEqual(model.model_name, DEFAULT_MODEL_NAME) -@unittest.skipIf(not _HF_TOKEN, 'HF_TOKEN environment variable not set.') +@unittest.skipIf( + not _HF_TOKEN, 'HF_INFERENCE_TOKEN environment variable not set.') class HuggingfaceInferenceAPITest(unittest.TestCase): def setUp(self): self.artifact_location = tempfile.mkdtemp() @@ -494,6 +522,9 @@ def test_get_api_url_and_when_model_name_not_provided(self): _ = inference_embeddings.api_url def test_embeddings_with_inference_api(self): + ready, reason = _hf_inference_api_available() + if not ready: + self.skipTest(reason) embedding_config = InferenceAPIEmbeddings( hf_token=_HF_TOKEN, model_name=INFERENCE_API_MODEL_NAME, @@ -516,7 +547,8 @@ def _embedding_shape(element): equal_to([(INFERENCE_API_EMBEDDING_DIM, True)])) -@unittest.skipIf(not _HF_TOKEN, 'HF_TOKEN environment variable not set.') +@unittest.skipIf( + not _HF_TOKEN, 'HF_INFERENCE_TOKEN environment variable not set.') class HuggingfaceInferenceAPIGCSLocationTest(HuggingfaceInferenceAPITest): def setUp(self): self.artifact_location = self.gcs_artifact_location = os.path.join(