[Yaml] - fix ml embeddings issue round2 (#35756)

* initial code

* rest of code with some improvements

* address comments and update design

* change tensorflow

* fix lint

* revert tensorflow version change

Author: derrickaw

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

@@ -183,20 +183,23 @@ 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() if hasattr(row, '_asdict') else row for row in batch]
+
   if not batch or not isinstance(batch[0], dict):
     raise TypeError(
         'Expected data to be dicts, got '
         f'{type(batch[0])} instead.')
-
   result = []
   expected_keys = set(batch[0].keys())
   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 +215,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() if hasattr(row, '_asdict') else row for row in batch]
+
   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,42 @@
 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 = []
+  skip_modules = ['base', 'handlers', 'test', 'tft', 'utils']
+  for _, module_name, _ in pkgutil.walk_packages(
+      package.__path__, package.__name__ + '.'):
+    if any(skip_name in module_name for skip_name in skip_modules):
+      continue
+    submodules.append(module_name)
+  return submodules
+
+
+_transform_constructors = {}
 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__
 except ImportError:
   tft = None  # type: ignore
 
+if tft:
+  # 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,68 @@ 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)))
+
+  def test_sentence_transformer_embedding_with_beam_rows(self):
+    SENTENCE_EMBEDDING_DIMENSION = 384
+    DATA = [
+        beam.Row(id=1, log_message="Error in module A"),
+        beam.Row(id=2, log_message="Warning in module B"),
+        beam.Row(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)