Support Embeddings in mltransform

sdks/python/apache_beam/ml/transforms/base.py

@@ -96,8 +96,8 @@ class ArtifactMode(object):
 class MLTransformProvider:
   """
   Data processing transforms that are intended to be used with MLTransform
-  should subclass MLTransformProvider and implement the following methods:
-  1. get_ptransform_for_processing()
+  should subclass MLTransformProvider and implement
+  get_ptransform_for_processing().
 
   get_ptransform_for_processing() method should return a PTransform that can be
   used to process the data.
@@ -184,7 +184,7 @@ def __init__(
     if kwargs:
       _LOGGER.warning("Ignoring the following arguments: %s", kwargs.keys())
 
-  # TODO: Add set_model_handler method.
+  # TODO:https://github.com/apache/beam/pull/29564 add set_model_handler method
   @abc.abstractmethod
   def get_model_handler(self) -> ModelHandler:
     """
@@ -398,6 +398,17 @@ def save_attributes(
       artifact_location,
       **kwargs,
   ):
+    # if an artifact location is present, instead of overwriting the
+    # existing file, raise an error since the same artifact location
+    # can be used by multiple beam jobs and this could result in undesired
+    # behavior.
+    if FileSystems.exists(FileSystems.join(artifact_location,
+                                           _ATTRIBUTE_FILE_NAME)):
+      raise FileExistsError(
+          "The artifact location %s already exists and contains %s. Please "
+          "specify a different location." %
+          (artifact_location, _ATTRIBUTE_FILE_NAME))
+
     if _JsonPickleTransformAttributeManager._is_remote_path(artifact_location):
       temp_dir = tempfile.mkdtemp()
       temp_json_file = os.path.join(temp_dir, _ATTRIBUTE_FILE_NAME)

sdks/python/apache_beam/ml/transforms/base_test.py

@@ -567,6 +567,19 @@ def test_with_gcs_location_with_none_options(self):
       self.attribute_manager.save_attributes(
           ptransform_list=[], artifact_location=path)
 
+  def test_with_same_local_artifact_location(self):
+    artifact_location = self.artifact_location
+    attribute_manager = base._JsonPickleTransformAttributeManager()
+
+    ptransform_list = [RunInference(model_handler=FakeModelHandler())]
+
+    attribute_manager.save_attributes(
+        ptransform_list, artifact_location=artifact_location)
+
+    with self.assertRaises(FileExistsError):
+      attribute_manager.save_attributes([lambda x: x],
+                                        artifact_location=artifact_location)
+
 
 if __name__ == '__main__':
   unittest.main()

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

@@ -14,12 +14,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import shutil
 import tempfile
 import unittest
+import uuid
 
 import numpy as np
-import pytest
 from parameterized import parameterized
 
 import apache_beam as beam
@@ -80,10 +81,11 @@
 @unittest.skipIf(
     SentenceTransformerEmbeddings is None,
     'sentence-transformers is not installed.')
-@pytest.mark.no_xdist
 class SentenceTrasformerEmbeddingsTest(unittest.TestCase):
   def setUp(self) -> None:
-    self.artifact_location = tempfile.mkdtemp()
+    self.artifact_location = tempfile.mkdtemp(prefix='sentence_transformers_')
+    self.gcs_artifact_location = os.path.join(
+        'gs://apache-beam-ml/testing/sentence_transformers', uuid.uuid4().hex)
 
   def tearDown(self) -> None:
     shutil.rmtree(self.artifact_location)
@@ -178,7 +180,6 @@ def test_sentence_transformer_with_int_data_types(self):
 
   @parameterized.expand(_parameterized_inputs)
   def test_with_gcs_artifact_location(self, inputs, model_name, output):
-    artifact_location = ('gs://apache-beam-ml/testing/sentence_transformers')
     embedding_config = SentenceTransformerEmbeddings(
         model_name=model_name, columns=[test_query_column])
 
@@ -187,8 +188,8 @@ def test_with_gcs_artifact_location(self, inputs, model_name, output):
           p
           | "CreateData" >> beam.Create(inputs)
           | "MLTransform" >>
-          MLTransform(write_artifact_location=artifact_location).with_transform(
-              embedding_config))
+          MLTransform(write_artifact_location=self.gcs_artifact_location
+                      ).with_transform(embedding_config))
       max_ele_pcoll = (
           result_pcoll
           | beam.Map(lambda x: round(np.max(x[test_query_column]), 2)))
@@ -200,7 +201,7 @@ def test_with_gcs_artifact_location(self, inputs, model_name, output):
           p
           | "CreateData" >> beam.Create(inputs)
           | "MLTransform" >>
-          MLTransform(read_artifact_location=artifact_location))
+          MLTransform(read_artifact_location=self.gcs_artifact_location))
       max_ele_pcoll = (
           result_pcoll
           | beam.Map(lambda x: round(np.max(x[test_query_column]), 2)))
@@ -233,7 +234,7 @@ def assert_element(element):
           | beam.Map(lambda x: x[test_query_column])
           | beam.Map(assert_element))
 
-  def test_mltransform_to_ptransform_with_vertex(self):
+  def test_mltransform_to_ptransform_with_sentence_transformer(self):
     model_name = ''
     transforms = [
         SentenceTransformerEmbeddings(columns=['x'], model_name=model_name),

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

@@ -118,9 +118,10 @@ def __init__(
     Args:
       model_name: The name of the Vertex AI Text Embedding model.
       columns: The columns containing the text to be embedded.
-      task_type: The downstream task for the embeddings.
-        Valid values are RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT,
-        SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING.
+      task_type: The downstream task for the embeddings. Valid values are
+        RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT, SEMANTIC_SIMILARITY,
+        CLASSIFICATION, CLUSTERING. For more information on the task type,
+        look at https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings # pylint: disable=line-too-long
       title: Identifier of the text content.
       project: The default GCP project for API calls.
       location: The default location for API calls.

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

@@ -14,9 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import shutil
 import tempfile
 import unittest
+import uuid
 
 import apache_beam as beam
 from apache_beam.ml.inference.base import RunInference
@@ -44,7 +46,9 @@
     VertexAITextEmbeddings is None, 'Vertex AI Python SDK is not installed.')
 class VertexAIEmbeddingsTest(unittest.TestCase):
   def setUp(self) -> None:
-    self.artifact_location = tempfile.mkdtemp()
+    self.artifact_location = tempfile.mkdtemp(prefix='_vertex_ai_test')
+    self.gcs_artifact_location = os.path.join(
+        'gs://apache-beam-ml/testing/vertex_ai', uuid.uuid4().hex)
 
   def tearDown(self) -> None:
     shutil.rmtree(self.artifact_location)
@@ -158,7 +162,6 @@ def test_with_int_data_types(self):
                     embedding_config))
 
   def test_with_gcs_artifact_location(self):
-    artifact_location = ('gs://apache-beam-ml/testing/vertex_ai')
     with beam.Pipeline() as p:
       embedding_config = VertexAITextEmbeddings(
           model_name=model_name, columns=[test_query_column])
@@ -172,7 +175,7 @@ def test_with_gcs_artifact_location(self):
         _ = self.pipeline_with_configurable_artifact_location(
             pipeline=data,
             embedding_config=embedding_config,
-            write_artifact_location=artifact_location)
+            write_artifact_location=self.gcs_artifact_location)
 
       with beam.Pipeline() as p:
         data = (
@@ -183,7 +186,7 @@ def test_with_gcs_artifact_location(self):
                 test_query_column: test_query
             }]))
         result_pcoll = self.pipeline_with_configurable_artifact_location(
-            pipeline=data, read_artifact_location=artifact_location)
+            pipeline=data, read_artifact_location=self.gcs_artifact_location)
 
         def assert_element(element):
           assert round(element, 2) == 0.15

sdks/python/apache_beam/ml/transforms/tft.py

@@ -105,10 +105,6 @@ def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform:
           "artifact_location is not specified. Please specify the "
           "artifact_location for the op %s" % self.__class__.__name__)
 
-    transforms = kwargs.get('transforms')
-    if transforms:
-      params['transforms'] = transforms
-
     artifact_mode = kwargs.get('artifact_mode')
     if artifact_mode:
       params['artifact_mode'] = artifact_mode

sdks/python/tox.ini

@@ -425,12 +425,13 @@ commands =
   /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_vertex_ai {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
 
 
-[testenv:py{38,39,310,311}-sentence-transformers-222]
+[testenv:py{38,39,310,311}-embeddings]
 deps =
   sentence-transformers==2.2.2
 extras = test,gcp
 commands =
   # Log aiplatform and its dependencies version for debugging
   /bin/sh -c "pip freeze | grep -E sentence-transformers"
+  /bin/sh -c "pip freeze | grep -E google-cloud-aiplatform"
   # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
-  /bin/sh -c 'pytest apache_beam/ml/transforms/embeddings -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
+  /bin/sh -c 'pytest apache_beam/ml/transforms/embeddings -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'