[yaml] Fix ml embeddings issue (#35715)

* add a few more libraries

* add more libraries and test

* update types

* fix gemeni review comments

* protect input fn

* update types

* fix lint issues

* more lint

* fix ml transform imports

* update warning and make list_submodules private

Author: derrickaw

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

@@ -183,9 +183,13 @@ def append_transform(self, transform: BaseOperation):
     """
 
 
-def _dict_input_fn(columns: Sequence[str],
-                   batch: Sequence[Dict[str, Any]]) -> List[str]:
+def _dict_input_fn(
+    columns: Sequence[str], batch: Sequence[Union[Dict[str, Any],
+                                                  beam.Row]]) -> List[str]:
   """Extract text from specified columns in batch."""
+  if batch and hasattr(batch[0], '_asdict'):
+    batch = [row._asdict() for row in batch if hasattr(row, '_asdict')]
+
   if not batch or not isinstance(batch[0], dict):
     raise TypeError(
         'Expected data to be dicts, got '
@@ -196,7 +200,7 @@ def _dict_input_fn(columns: Sequence[str],
   expected_columns = set(columns)
   # Process one batch item at a time
   for item in batch:
-    item_keys = item.keys()
+    item_keys = item.keys() if isinstance(item, dict) else set()
     if set(item_keys) != expected_keys:
       extra_keys = item_keys - expected_keys
       missing_keys = expected_keys - item_keys
@@ -212,21 +216,31 @@ def _dict_input_fn(columns: Sequence[str],
 
     # Get all columns for this item
     for col in columns:
-      result.append(item[col])
+      if isinstance(item, dict):
+        result.append(item[col])
   return result
 
 
 def _dict_output_fn(
     columns: Sequence[str],
-    batch: Sequence[Dict[str, Any]],
-    embeddings: Sequence[Any]) -> List[Dict[str, Any]]:
+    batch: Sequence[Union[Dict[str, Any], beam.Row]],
+    embeddings: Sequence[Any]) -> list[Union[dict[str, Any], beam.Row]]:
   """Map embeddings back to columns in batch."""
+  is_beam_row = False
+  if batch and hasattr(batch[0], '_asdict'):
+    is_beam_row = True
+    batch = [row._asdict() for row in batch if hasattr(row, '_asdict')]
+
   result = []
   for batch_idx, item in enumerate(batch):
     for col_idx, col in enumerate(columns):
       embedding_idx = batch_idx * len(columns) + col_idx
-      item[col] = embeddings[embedding_idx]
+      if isinstance(item, dict):
+        item[col] = embeddings[embedding_idx]
     result.append(item)
+
+  if is_beam_row:
+    result = [beam.Row(**item) for item in result if isinstance(item, dict)]
   return result
 
 

sdks/python/apache_beam/yaml/yaml_ml.py

@@ -16,7 +16,10 @@
 #
 
 """This module defines yaml wrappings for some ML transforms."""
+import logging
+import pkgutil
 from collections.abc import Callable
+from importlib import import_module
 from typing import Any
 from typing import Optional
 
@@ -29,14 +32,37 @@
 from apache_beam.yaml import options
 from apache_beam.yaml.yaml_utils import SafeLineLoader
 
+
+def _list_submodules(package):
+  """
+    Lists all submodules within a given package.
+    """
+  submodules = []
+  for _, module_name, _ in pkgutil.walk_packages(
+      package.__path__, package.__name__ + '.'):
+    if 'test' in module_name:
+      continue
+    submodules.append(module_name)
+  return submodules
+
+
 try:
   from apache_beam.ml.transforms import tft
   from apache_beam.ml.transforms.base import MLTransform
   # TODO(robertwb): Is this all of them?
-  _transform_constructors = tft.__dict__
+  _transform_constructors = {}
 except ImportError:
   tft = None  # type: ignore
 
+# Load all available ML Transform modules
+for module_name in _list_submodules(beam.ml.transforms):
+  try:
+    module = import_module(module_name)
+    _transform_constructors |= module.__dict__
+  except ImportError as e:
+    logging.warning('Could not load ML transform module %s: %s.  Please ' \
+                    'install the necessary module dependencies', module_name, e)
+
 
 class ModelHandlerProvider:
   handler_types: dict[str, Callable[..., "ModelHandlerProvider"]] = {}

sdks/python/apache_beam/yaml/yaml_ml_test.py

@@ -86,6 +86,38 @@ def test_ml_transform(self):
             equal_to([5]),
             label='CheckVocab')
 
+  def test_sentence_transformer_embedding(self):
+    SENTENCE_EMBEDDING_DIMENSION = 384
+    DATA = [{
+        'id': 1, 'log_message': "Error in module A"
+    }, {
+        'id': 2, 'log_message': "Warning in module B"
+    }, {
+        'id': 3, 'log_message': "Info in module C"
+    }]
+    ml_opts = beam.options.pipeline_options.PipelineOptions(
+        pickle_library='cloudpickle', yaml_experimental_features=['ML'])
+    with tempfile.TemporaryDirectory() as tempdir:
+      with beam.Pipeline(options=ml_opts) as p:
+        elements = p | beam.Create(DATA)
+        result = elements | YamlTransform(
+            f'''
+            type: MLTransform
+            config:
+              write_artifact_location: {tempdir}
+              transforms:
+                - type: SentenceTransformerEmbeddings
+                  config:
+                    model_name: all-MiniLM-L6-v2
+                    columns: [log_message]
+            ''')
+
+        # Perform a basic check to ensure that embeddings are generated
+        # and that the dimension of those embeddings is correct.
+        actual_output = result | beam.Map(lambda x: len(x['log_message']))
+        assert_that(
+            actual_output, equal_to([SENTENCE_EMBEDDING_DIMENSION] * len(DATA)))
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)

sdks/python/setup.py

@@ -518,6 +518,7 @@ def get_portability_package_data():
               'pyod',
               'tensorflow',
               'tensorflow-hub',
+              'tensorflow-text',
               'tensorflow-transform',
               'tf2onnx',
               'torch',