(beam) branch master updated: Fix HF Inference API test model (#40844)

[yhu]

This is an automated email from the ASF dual-hosted git repository.

Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 688a9c11da7 Fix HF Inference API test model (#40844)
688a9c11da7 is described below

commit 688a9c11da70490ca93fa297f34fb4a31f5f4cfe
Author: Abdelrahman Ibrahim <abdoibrahim1...@gmail.com>
AuthorDate: Wed Oct 28 17:49:50 2026 +0300

    Fix HF Inference API test model (#40844)
    
    * fix HF Inference API test model
    
    * trigger workflow
    
    * trigger workflow
---
 .../beam_PostCommit_Python_Dependency.json             |  2 +-
 .../ml/transforms/embeddings/huggingface_test.py       | 18 ++++++++++++------
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json 
b/.github/trigger_files/beam_PostCommit_Python_Dependency.json
index 0c0e8ca6a12..a5413bd5db9 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": 10
   }
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 12990cb4c2c..166827f3e90 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 @@ _HF_TOKEN = os.environ.get('HF_INFERENCE_TOKEN')
 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 @@ class HuggingfaceInferenceAPITest(unittest.TestCase):
   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 @@ class HuggingfaceInferenceAPITest(unittest.TestCase):
           | "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.')