From 25f71b0ab24566500507c5a4dbb840974c8a05da Mon Sep 17 00:00:00 2001 From: aibrahiim Date: Tue, 27 Oct 2026 20:38:54 +0300 Subject: [PATCH 1/3] fix HF Inference API test model --- .../transforms/embeddings/huggingface_test.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) 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 12990cb4c2c8..166827f3e907 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py @@ -66,6 +66,10 @@ test_query = "This is a test" test_query_column = "feature_1" DEFAULT_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2" +# all-mpnet-base-v2 is only served as sentence-similarity on hf-inference; +# InferenceAPIEmbeddings calls /pipeline/feature-extraction. +INFERENCE_API_MODEL_NAME = "BAAI/bge-small-en-v1.5" +INFERENCE_API_EMBEDDING_DIM = 384 IMAGE_MODEL_NAME = "clip-ViT-B-32" _parameterized_inputs = [ ([{ @@ -492,10 +496,9 @@ def test_get_api_url_and_when_model_name_not_provided(self): def test_embeddings_with_inference_api(self): embedding_config = InferenceAPIEmbeddings( hf_token=_HF_TOKEN, - model_name=DEFAULT_MODEL_NAME, + model_name=INFERENCE_API_MODEL_NAME, columns=[test_query_column], ) - expected_output = [0.13] with beam.Pipeline() as p: result_pcoll = ( p @@ -503,11 +506,14 @@ def test_embeddings_with_inference_api(self): | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) - max_ele_pcoll = ( - result_pcoll - | beam.Map(lambda x: round(np.max(x[test_query_column]), 2))) - assert_that(max_ele_pcoll, equal_to(expected_output)) + def _embedding_shape(element): + emb = element[test_query_column] + return (len(emb), all(isinstance(x, (float, np.floating)) for x in emb)) + + assert_that( + result_pcoll | beam.Map(_embedding_shape), + equal_to([(INFERENCE_API_EMBEDDING_DIM, True)])) @unittest.skipIf(not _HF_TOKEN, 'HF_TOKEN environment variable not set.') From 4a1e2d9c835e72fb31856c3fdb5751eb02f09be5 Mon Sep 17 00:00:00 2001 From: aibrahiim Date: Tue, 27 Oct 2026 20:41:22 +0300 Subject: [PATCH 2/3] trigger workflow --- .github/trigger_files/beam_PostCommit_Python_Dependency.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json b/.github/trigger_files/beam_PostCommit_Python_Dependency.json index 0c0e8ca6a12b..f3bc9e958674 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": 8 + "modification": 9 } From dc2577019f4a29fed920842e002d30fe2fdf6372 Mon Sep 17 00:00:00 2001 From: aibrahiim Date: Tue, 27 Oct 2026 23:03:25 +0300 Subject: [PATCH 3/3] trigger workflow --- .github/trigger_files/beam_PostCommit_Python_Dependency.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json b/.github/trigger_files/beam_PostCommit_Python_Dependency.json index f3bc9e958674..a5413bd5db93 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": 10 }