diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json b/.github/trigger_files/beam_PostCommit_Python_Dependency.json index 0c0e8ca6a12b..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": 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 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.')