Removes usages of the sentence-transformers package...

sdks/python/apache_beam/ml/rag/embeddings/huggingface.py

@@ -32,11 +32,6 @@
 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 @@ def __init__(
               ``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,24 +155,19 @@ def __init__(
               ``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. "
           "Please install it with `pip install pillow`.")
     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,

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 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 @@ def test_image_embedding_pipeline_from_path(self):
       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

sdks/python/apache_beam/ml/transforms/embeddings/huggingface.py

@@ -21,6 +21,7 @@
 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 _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,30 +77,35 @@ def __init__(
       **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
     self._max_batch_size = max_batch_size
     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 @@ def __init__(
     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,