This is an automated email from the ASF dual-hosted git repository.
chamikaramj 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 fadfc789376 Removes usages of the sentence-transformers package...
(#40439)
fadfc789376 is described below
commit fadfc789376137afabab0469bf1e43a0f10a6ad0
Author: Chamikara Jayalath <[email protected]>
AuthorDate: Thu Oct 8 19:03:09 2026 -0700
Removes usages of the sentence-transformers package... (#40439)
---
.../apache_beam/ml/rag/embeddings/huggingface.py | 23 +---
.../ml/rag/embeddings/huggingface_test.py | 56 ++++++++
.../ml/transforms/embeddings/huggingface.py | 48 +++++--
.../ml/transforms/embeddings/huggingface_test.py | 142 +++++++++++++++++++++
sdks/python/apache_beam/testing/test_utils.py | 61 +++++++++
sdks/python/apache_beam/testing/test_utils_test.py | 31 +++++
sdks/python/apache_beam/yaml/yaml_ml_test.py | 40 +++++-
7 files changed, 368 insertions(+), 33 deletions(-)
diff --git a/sdks/python/apache_beam/ml/rag/embeddings/huggingface.py
b/sdks/python/apache_beam/ml/rag/embeddings/huggingface.py
index e17a5d66fe9..8ad34ca690b 100644
--- a/sdks/python/apache_beam/ml/rag/embeddings/huggingface.py
+++ b/sdks/python/apache_beam/ml/rag/embeddings/huggingface.py
@@ -32,11 +32,6 @@ from apache_beam.ml.transforms.base import
_ImageEmbeddingHandler
from apache_beam.ml.transforms.base import _TextEmbeddingHandler
from apache_beam.ml.transforms.embeddings.huggingface import
_SentenceTransformerModelHandler
-try:
- from sentence_transformers import SentenceTransformer
-except ImportError:
- SentenceTransformer = None
-
try:
from PIL import Image as PILImage
except ImportError:
@@ -66,20 +61,15 @@ class HuggingfaceTextEmbeddings(EmbeddingsManager):
``model.encode()``
(e.g. ``normalize_embeddings``).
"""
- if not SentenceTransformer:
- raise ImportError(
- "sentence-transformers is required to use "
- "HuggingfaceTextEmbeddings."
- "Please install it with using `pip install sentence-transformers`.")
super().__init__(type_adapter=create_text_adapter(), **kwargs)
self.model_name = model_name
self.max_seq_length = max_seq_length
- self.model_class = SentenceTransformer
def get_model_handler(self):
"""Returns model handler configured with RAG adapter."""
+ # model_class is not set, so that sentence-transformers is only imported
+ # when the model is loaded during pipeline execution.
return _SentenceTransformerModelHandler(
- model_class=self.model_class,
max_seq_length=self.max_seq_length,
model_name=self.model_name,
load_model_args=self.load_model_args,
@@ -165,11 +155,6 @@ class HuggingfaceImageEmbeddings(EmbeddingsManager):
``model.encode()``
(e.g. ``normalize_embeddings``).
"""
- if not SentenceTransformer:
- raise ImportError(
- "sentence-transformers is required to use "
- "HuggingfaceImageEmbeddings. "
- "Please install it with `pip install sentence-transformers`.")
if not PILImage:
raise ImportError(
"Pillow is required to use HuggingfaceImageEmbeddings. "
@@ -177,12 +162,12 @@ class HuggingfaceImageEmbeddings(EmbeddingsManager):
super().__init__(type_adapter=_create_hf_image_adapter(), **kwargs)
self.model_name = model_name
self.max_seq_length = max_seq_length
- self.model_class = SentenceTransformer
def get_model_handler(self):
"""Returns model handler configured with RAG adapter."""
+ # model_class is not set, so that sentence-transformers is only imported
+ # when the model is loaded during pipeline execution.
return _SentenceTransformerModelHandler(
- model_class=self.model_class,
max_seq_length=self.max_seq_length,
model_name=self.model_name,
load_model_args=self.load_model_args,
diff --git a/sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py
b/sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py
index 65c21771524..de93df95627 100644
--- a/sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py
+++ b/sdks/python/apache_beam/ml/rag/embeddings/huggingface_test.py
@@ -16,6 +16,7 @@
"""Tests for apache_beam.ml.rag.embeddings.huggingface."""
+import importlib.util
import io
import os
import shutil
@@ -33,6 +34,7 @@ from apache_beam.ml.rag.types import Content
from apache_beam.ml.rag.types import EmbeddableItem
from apache_beam.ml.rag.types import Embedding
from apache_beam.ml.transforms.base import MLTransform
+from apache_beam.testing import test_utils
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
@@ -234,6 +236,60 @@ class HuggingfaceImageEmbeddingsTest(unittest.TestCase):
shutil.rmtree(artifact_location)
+class HuggingfaceEmbeddingsJobSubmissionTest(unittest.TestCase):
+ """Tests that sentence-transformers is only used when executing pipelines.
+
+ sentence-transformers should only be needed for loading the model when the
+ pipeline is executed (for example, on the workers), but not for constructing
+ and submitting the pipeline. These tests don't require sentence-transformers
+ to be installed.
+ """
+ def setUp(self):
+ self.artifact_location = tempfile.mkdtemp(prefix='hf_job_submission_')
+
+ def tearDown(self):
+ shutil.rmtree(self.artifact_location)
+
+ def _construct_and_serialize_pipeline(self, items, embedder):
+ pipeline = beam.Pipeline()
+ _ = (
+ pipeline
+ | beam.Create(items)
+ | MLTransform(write_artifact_location=self.artifact_location).
+ with_transform(embedder))
+ # Serializes the pipeline, as done during job submission, and deserializes
+ # it again. Deserializing fails if the pipeline references (i.e. requires)
+ # a module that can't be imported.
+ beam.Pipeline.from_runner_api(
+ pipeline.to_runner_api(), pipeline.runner, pipeline.options)
+
+ def test_importing_module_does_not_import_sentence_transformers(self):
+ # Executes the module in a new module object, i.e. without replacing the
+ # already imported module.
+ spec = importlib.util.find_spec(HuggingfaceTextEmbeddings.__module__)
+ with test_utils.block_imports('sentence_transformers') as
attempted_imports:
+ spec.loader.exec_module(importlib.util.module_from_spec(spec))
+ self.assertEqual(attempted_imports, [])
+
+ def test_text_embeddings_job_submission_does_not_use_sentence_transformers(
+ self):
+ items = [Chunk(content=Content(text="This is a test sentence."), id="1")]
+ with test_utils.block_imports('sentence_transformers') as
attempted_imports:
+ embedder = HuggingfaceTextEmbeddings(
+ model_name="sentence-transformers/all-MiniLM-L6-v2")
+ self._construct_and_serialize_pipeline(items, embedder)
+ self.assertEqual(attempted_imports, [])
+
+ @unittest.skipIf(not PIL_AVAILABLE, "PIL not available")
+ def test_image_embeddings_job_submission_does_not_use_sentence_transformers(
+ self):
+ items = [EmbeddableItem.from_image(_create_png_bytes(), id="img1")]
+ with test_utils.block_imports('sentence_transformers') as
attempted_imports:
+ embedder = HuggingfaceImageEmbeddings(model_name="clip-ViT-B-32")
+ self._construct_and_serialize_pipeline(items, embedder)
+ self.assertEqual(attempted_imports, [])
+
+
def _create_png_bytes():
"""Create a small valid RGB PNG image."""
from PIL import Image
diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py
b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py
index a6905c37545..c3eeccf2a0f 100644
--- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py
+++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py
@@ -21,6 +21,7 @@ import os
from collections.abc import Callable
from collections.abc import Mapping
from collections.abc import Sequence
+from typing import TYPE_CHECKING
from typing import Any
from typing import Optional
@@ -33,14 +34,29 @@ from apache_beam.ml.transforms.base import EmbeddingsManager
from apache_beam.ml.transforms.base import _ImageEmbeddingHandler
from apache_beam.ml.transforms.base import _TextEmbeddingHandler
-try:
+# sentence-transformers is intentionally not imported at module import time.
+# It is only needed for loading the model when the pipeline is executed (for
+# example, on the workers) and should not be required for constructing and
+# submitting pipelines.
+if TYPE_CHECKING:
from sentence_transformers import SentenceTransformer
-except ImportError:
- SentenceTransformer = None
_LOGGER = logging.getLogger(__name__)
+def _import_sentence_transformer() -> Callable:
+ """Returns sentence_transformers.SentenceTransformer, importing it lazily."""
+ try:
+ from sentence_transformers import SentenceTransformer
+ except ImportError as e:
+ raise ImportError(
+ "sentence-transformers is required to load SentenceTransformer "
+ "models, e.g. for SentenceTransformerEmbeddings. Please install it "
+ "using `pip install sentence-transformers` in the environment where "
+ "the pipeline is executed (for example, on the workers).") from e
+ return SentenceTransformer
+
+
# TODO: https://github.com/apache/beam/issues/29621
# Use HuggingFaceModelHandlerTensor once the import issue is fixed.
# Right now, the hugging face model handler import torch and tensorflow
@@ -52,7 +68,7 @@ class _SentenceTransformerModelHandler(ModelHandler):
def __init__(
self,
model_name: str,
- model_class: Callable,
+ model_class: Optional[Callable] = None,
load_model_args: Optional[dict] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
@@ -61,6 +77,9 @@ class _SentenceTransformerModelHandler(ModelHandler):
**kwargs):
self._max_seq_length = max_seq_length
self.model_name = model_name
+ # If not set, sentence_transformers.SentenceTransformer is used. It is
+ # imported in load_model(), i.e. only when the pipeline is executed, so
that
+ # sentence-transformers is not required for job submission.
self._model_class = model_class
self._load_model_args = load_model_args
self._min_batch_size = min_batch_size
@@ -68,23 +87,25 @@ class _SentenceTransformerModelHandler(ModelHandler):
self._large_model = large_model
self._kwargs = kwargs
- if not SentenceTransformer:
- raise ImportError(
- "sentence-transformers is required to use "
- "SentenceTransformerEmbeddings."
- "Please install it with using `pip install sentence-transformers`.")
-
def run_inference(
self,
batch: Sequence[str],
- model: SentenceTransformer,
+ model: 'SentenceTransformer',
inference_args: Optional[dict[str, Any]] = None,
):
inference_args = inference_args or {}
return model.encode(batch, **inference_args)
def load_model(self):
- model = self._model_class(self.model_name, **self._load_model_args)
+ model_class = self._model_class
+ # Checking callable() instead of None also covers artifacts written by
+ # Beam 2.77.0 and earlier, which stored a reference to the
+ # SentenceTransformer class. jsonpickle restores such a reference as a
+ # non-callable placeholder if sentence-transformers was not installed when
+ # the artifacts were read (for example, during job submission).
+ if not callable(model_class):
+ model_class = _import_sentence_transformer()
+ model = model_class(self.model_name, **(self._load_model_args or {}))
if self._max_seq_length:
model.max_seq_length = self._max_seq_length
return model
@@ -139,8 +160,9 @@ class SentenceTransformerEmbeddings(EmbeddingsManager):
self.image_model = image_model
def get_model_handler(self):
+ # model_class is not set, so that sentence-transformers is only imported
+ # when the model is loaded during pipeline execution.
return _SentenceTransformerModelHandler(
- model_class=SentenceTransformer,
max_seq_length=self.max_seq_length,
model_name=self.model_name,
load_model_args=self.load_model_args,
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 5e56dc71054..12990cb4c2c 100644
--- a/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py
+++ b/sdks/python/apache_beam/ml/transforms/embeddings/huggingface_test.py
@@ -14,12 +14,17 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import contextlib
+import importlib.util
import os
import shutil
+import sys
import tempfile
+import types
import unittest
import uuid
+import jsonpickle
import numpy as np
import pytest
from parameterized import parameterized
@@ -28,6 +33,8 @@ import apache_beam as beam
from apache_beam.ml.inference.base import RunInference
from apache_beam.ml.transforms import base
from apache_beam.ml.transforms.base import MLTransform
+from apache_beam.ml.transforms.embeddings import huggingface
+from apache_beam.testing import test_utils
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
@@ -330,6 +337,141 @@ class
SentenceTransformerEmbeddingsTest(unittest.TestCase):
embedding_config))
+class _FakeSentenceTransformer:
+ """A stand-in for sentence_transformers.SentenceTransformer."""
+ def __init__(self, model_name, **kwargs):
+ self.model_name = model_name
+ self.kwargs = kwargs
+ self.max_seq_length = None
+
+
[email protected]
+def _fake_sentence_transformers_module():
+ """Replaces the sentence_transformers module with a fake one."""
+ fake_module = types.ModuleType('sentence_transformers')
+ fake_module.SentenceTransformer = _FakeSentenceTransformer
+ # Also hides the real sentence_transformers modules, if already imported, and
+ # restores them afterwards.
+ with test_utils.block_imports('sentence_transformers'):
+ sys.modules['sentence_transformers'] = fake_module
+ try:
+ yield
+ finally:
+ del sys.modules['sentence_transformers']
+
+
+def _serialize_and_deserialize(pipeline):
+ """Serializes the pipeline, as done during job submission, and deserializes
+ it again. Deserializing fails if the pipeline references (i.e. requires) a
+ module that can't be imported."""
+ return beam.Pipeline.from_runner_api(
+ pipeline.to_runner_api(), pipeline.runner, pipeline.options)
+
+
+class SentenceTransformerJobSubmissionTest(unittest.TestCase):
+ """Tests that sentence-transformers is only used when executing pipelines.
+
+ sentence-transformers (and its heavy dependencies, like torch) should only be
+ needed for loading the model when the pipeline is executed (for example, on
+ the workers), but not for constructing and submitting the pipeline. These
+ tests don't require sentence-transformers to be installed.
+ """
+ def setUp(self):
+ self.artifact_location = tempfile.mkdtemp(prefix='sentence_transformers_')
+
+ def tearDown(self):
+ shutil.rmtree(self.artifact_location)
+
+ def _apply_ml_transform(self, pipeline, ml_transform):
+ _ = (
+ pipeline
+ | beam.Create([{
+ test_query_column: test_query
+ }])
+ | ml_transform)
+
+ def test_importing_module_does_not_import_sentence_transformers(self):
+ # Executes the module in a new module object, i.e. without replacing the
+ # already imported module.
+ spec = importlib.util.find_spec(huggingface.__name__)
+ with test_utils.block_imports('sentence_transformers') as
attempted_imports:
+ spec.loader.exec_module(importlib.util.module_from_spec(spec))
+ self.assertEqual(attempted_imports, [])
+
+ @parameterized.expand([('text_model', False), ('image_model', True)])
+ def test_job_submission_does_not_use_sentence_transformers(
+ self, unused_name, image_model):
+ with test_utils.block_imports('sentence_transformers') as
attempted_imports:
+ embedding_config = huggingface.SentenceTransformerEmbeddings(
+ model_name=DEFAULT_MODEL_NAME,
+ columns=[test_query_column],
+ image_model=image_model)
+ pipeline = beam.Pipeline()
+ self._apply_ml_transform(
+ pipeline,
+ MLTransform(
+ write_artifact_location=self.artifact_location).with_transform(
+ embedding_config))
+ _serialize_and_deserialize(pipeline)
+ self.assertEqual(attempted_imports, [])
+
+ def test_reading_artifacts_does_not_use_sentence_transformers(self):
+ with test_utils.block_imports('sentence_transformers') as
attempted_imports:
+ embedding_config = huggingface.SentenceTransformerEmbeddings(
+ model_name=DEFAULT_MODEL_NAME, columns=[test_query_column])
+ # Artifacts are written when constructing the pipeline.
+ self._apply_ml_transform(
+ beam.Pipeline(),
+ MLTransform(
+ write_artifact_location=self.artifact_location).with_transform(
+ embedding_config))
+ pipeline = beam.Pipeline()
+ self._apply_ml_transform(
+ pipeline, MLTransform(read_artifact_location=self.artifact_location))
+ _serialize_and_deserialize(pipeline)
+ self.assertEqual(attempted_imports, [])
+
+ def test_loading_model_requires_sentence_transformers(self):
+ model_handler = huggingface.SentenceTransformerEmbeddings(
+ model_name=DEFAULT_MODEL_NAME,
+ columns=[test_query_column]).get_model_handler()
+ with test_utils.block_imports('sentence_transformers'):
+ with self.assertRaisesRegex(ImportError,
+ 'sentence-transformers is required'):
+ model_handler.load_model()
+
+ def test_loading_model_uses_sentence_transformers(self):
+ model_handler = huggingface.SentenceTransformerEmbeddings(
+ model_name=DEFAULT_MODEL_NAME,
+ columns=[test_query_column],
+ max_seq_length=128,
+ load_model_args={
+ 'device': 'cpu'
+ }).get_model_handler()
+ with _fake_sentence_transformers_module():
+ model = model_handler.load_model()
+ self.assertIsInstance(model, _FakeSentenceTransformer)
+ self.assertEqual(model.model_name, DEFAULT_MODEL_NAME)
+ self.assertEqual(model.kwargs, {'device': 'cpu'})
+ self.assertEqual(model.max_seq_length, 128)
+
+ def test_loading_model_from_artifacts_written_by_older_versions(self):
+ model_handler = huggingface.SentenceTransformerEmbeddings(
+ model_name=DEFAULT_MODEL_NAME,
+ columns=[test_query_column]).get_model_handler()
+ # Beam 2.77.0 and earlier stored a reference to the SentenceTransformer
+ # class in the artifacts. This is how jsonpickle restores the reference if
+ # sentence-transformers is not installed when reading the artifacts.
+ with test_utils.block_imports('sentence_transformers'):
+ model_handler._model_class = jsonpickle.decode(
+ '{"py/type": '
+ '"sentence_transformers.SentenceTransformer.SentenceTransformer"}')
+ with _fake_sentence_transformers_module():
+ model = model_handler.load_model()
+ self.assertIsInstance(model, _FakeSentenceTransformer)
+ self.assertEqual(model.model_name, DEFAULT_MODEL_NAME)
+
+
@unittest.skipIf(not _HF_TOKEN, 'HF_TOKEN environment variable not set.')
class HuggingfaceInferenceAPITest(unittest.TestCase):
def setUp(self):
diff --git a/sdks/python/apache_beam/testing/test_utils.py
b/sdks/python/apache_beam/testing/test_utils.py
index e31c53765b8..425b4c6dfac 100644
--- a/sdks/python/apache_beam/testing/test_utils.py
+++ b/sdks/python/apache_beam/testing/test_utils.py
@@ -22,11 +22,14 @@ For internal use only; no backwards-compatibility
guarantees.
# pytype: skip-file
+import contextlib
import hashlib
import importlib
+import importlib.abc
import os
import shutil
import struct
+import sys
import tempfile
import numpy as np
@@ -125,6 +128,64 @@ def patch_retry(testcase, module):
testcase.addCleanup(remove_patches)
+class _ImportBlocker(importlib.abc.MetaPathFinder):
+ """A meta path finder that blocks, and records, imports of given packages."""
+ def __init__(self, packages):
+ self._packages = tuple(packages)
+ self.attempted_imports = []
+
+ def blocks(self, module_name):
+ return any(
+ module_name == package or module_name.startswith(package + '.')
+ for package in self._packages)
+
+ def find_spec(self, fullname, path, target=None):
+ if self.blocks(fullname):
+ self.attempted_imports.append(fullname)
+ raise ModuleNotFoundError(
+ 'Import of %s is blocked for testing.' % fullname, name=fullname)
+ return None
+
+
[email protected]
+def block_imports(*packages):
+ """Context manager that makes the given packages unimportable.
+
+ Within the context, importing any of the given packages, or their
+ submodules, raises a ModuleNotFoundError as if they were not installed. This
+ also applies to packages that were already imported. Such packages are
+ restored when exiting the context.
+
+ This is useful, for example, for verifying that constructing and submitting
+ a pipeline neither requires, nor tries to import, a dependency that is only
+ needed when executing the pipeline on workers::
+
+ with block_imports('sentence_transformers') as attempted_imports:
+ ... # Construct and serialize the pipeline.
+ assert not attempted_imports
+
+ Args:
+ *packages: names of the top-level packages, or modules, to block.
+
+ Yields:
+ A list with the names of all modules whose import was attempted, and
+ blocked, within the context.
+ """
+ blocker = _ImportBlocker(packages)
+ hidden_modules = {
+ name: module
+ for name, module in list(sys.modules.items()) if blocker.blocks(name)
+ }
+ for name in hidden_modules:
+ del sys.modules[name]
+ sys.meta_path.insert(0, blocker)
+ try:
+ yield blocker.attempted_imports
+ finally:
+ sys.meta_path.remove(blocker)
+ sys.modules.update(hidden_modules)
+
+
@retry.with_exponential_backoff(
num_retries=3, retry_filter=retry.retry_on_beam_io_error_filter)
def delete_files(file_paths):
diff --git a/sdks/python/apache_beam/testing/test_utils_test.py
b/sdks/python/apache_beam/testing/test_utils_test.py
index 1a89ada2a70..9c379cf35aa 100644
--- a/sdks/python/apache_beam/testing/test_utils_test.py
+++ b/sdks/python/apache_beam/testing/test_utils_test.py
@@ -19,8 +19,10 @@
# pytype: skip-file
+import importlib
import logging
import os
+import sys
import tempfile
import unittest
@@ -95,6 +97,35 @@ class TestUtilsTest(unittest.TestCase):
pub_client.delete_topic.assert_called_with(topic=topic.name)
+class BlockImportsTest(unittest.TestCase):
+ def test_block_imports_blocks_and_records_imports(self):
+ with utils.block_imports('colorsys') as attempted_imports:
+ with self.assertRaises(ImportError):
+ importlib.import_module('colorsys')
+ self.assertEqual(attempted_imports, ['colorsys'])
+ # The module can be imported again after exiting the context.
+ importlib.import_module('colorsys')
+
+ def test_block_imports_hides_and_restores_imported_modules(self):
+ wsgiref_util = importlib.import_module('wsgiref.util')
+ wsgiref = sys.modules['wsgiref']
+ with utils.block_imports('wsgiref') as attempted_imports:
+ self.assertNotIn('wsgiref', sys.modules)
+ self.assertNotIn('wsgiref.util', sys.modules)
+ with self.assertRaises(ImportError):
+ importlib.import_module('wsgiref.util')
+ # Importing a submodule first imports its parent package.
+ self.assertEqual(attempted_imports, ['wsgiref'])
+ self.assertIs(sys.modules['wsgiref'], wsgiref)
+ self.assertIs(sys.modules['wsgiref.util'], wsgiref_util)
+
+ def test_block_imports_only_blocks_given_packages(self):
+ with utils.block_imports('wsgi', 'colorsys.foo') as attempted_imports:
+ importlib.import_module('wsgiref')
+ importlib.import_module('colorsys')
+ self.assertEqual(attempted_imports, [])
+
+
class LCGeneratorTest(unittest.TestCase):
Generator = utils.LCGenerator
diff --git a/sdks/python/apache_beam/yaml/yaml_ml_test.py
b/sdks/python/apache_beam/yaml/yaml_ml_test.py
index d8b1bdbae1b..c4fc179b9a3 100644
--- a/sdks/python/apache_beam/yaml/yaml_ml_test.py
+++ b/sdks/python/apache_beam/yaml/yaml_ml_test.py
@@ -20,6 +20,7 @@ import tempfile
import unittest
import apache_beam as beam
+from apache_beam.testing import test_utils
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.yaml.yaml_transform import YamlTransform
@@ -28,7 +29,7 @@ try:
# pylint: disable=wrong-import-order, wrong-import-position, unused-import
from apache_beam.ml.transforms import tft
except ImportError:
- raise unittest.SkipTest('tensorflow_transform is not installed.')
+ tft = None
TRAIN_DATA = [
beam.Row(num=0, text='And God said, Let there be light,'),
@@ -41,6 +42,7 @@ TEST_DATA = [
]
[email protected](tft is None, 'tensorflow_transform is not installed.')
class MLTransformTest(unittest.TestCase):
def test_ml_transform(self):
ml_opts = beam.options.pipeline_options.PipelineOptions(
@@ -236,6 +238,42 @@ class MLTransformTest(unittest.TestCase):
assert_that(result | beam.Map(check_row), equal_to([1, 2, 3]))
+class MLTransformJobSubmissionTest(unittest.TestCase):
+ def test_sentence_transformer_embedding_does_not_use_sentence_transformers(
+ self):
+ # sentence-transformers should only be needed when executing the pipeline
+ # (for example, on the workers), not for constructing and submitting it.
+ ml_opts = beam.options.pipeline_options.PipelineOptions(
+ pickle_library='cloudpickle', yaml_experimental_features=['ML'])
+ with tempfile.TemporaryDirectory() as tempdir:
+ with test_utils.block_imports(
+ 'sentence_transformers') as attempted_imports:
+ pipeline = beam.Pipeline(options=ml_opts)
+ _ = pipeline | YamlTransform(
+ f'''
+ type: chain
+ transforms:
+ - type: Create
+ config:
+ elements:
+ - {{id: 1, log_message: "Error in module A"}}
+ - type: MLTransform
+ config:
+ write_artifact_location: {tempdir}
+ transforms:
+ - type: SentenceTransformerEmbeddings
+ config:
+ model_name: all-MiniLM-L6-v2
+ columns: [log_message]
+ ''')
+ # Serializes the pipeline, as done during job submission, and
+ # deserializes it again. Deserializing fails if the pipeline references
+ # (i.e. requires) a module that can't be imported.
+ beam.Pipeline.from_runner_api(
+ pipeline.to_runner_api(), pipeline.runner, pipeline.options)
+ self.assertEqual(attempted_imports, [])
+
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
unittest.main()
