embeddings: use scale test validation

Author: moawnallah

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

@@ -24,14 +24,12 @@
 from apache_beam.ml.transforms import base
 from apache_beam.ml.transforms.base import MLTransform
 
-try:
-  from sdks.python.apache_beam.ml.transforms.embeddings.open_ai import OpenAITextEmbeddings
-except ImportError:
-  OpenAITextEmbeddings = None
+from apache_beam.ml.transforms.embeddings.open_ai import OpenAITextEmbeddings
 
 # pylint: disable=ungrouped-imports
 try:
   import tensorflow_transform as tft
+  from apache_beam.ml.transforms.tft import ScaleTo01
 except ImportError:
   tft = None
 
@@ -76,6 +74,7 @@ def test_embeddings_with_scale_to_0_1(self):
         columns=[test_query_column],
         api_key=self.api_key,
     )
+    scale_config = ScaleTo01(columns=['embedding'])
     with beam.Pipeline() as pipeline:
       transformed_pcoll = (
           pipeline
@@ -84,10 +83,12 @@ def test_embeddings_with_scale_to_0_1(self):
           }])
           | "MLTransform" >> MLTransform(
               write_artifact_location=self.artifact_location).with_transform(
-                  embedding_config))
+                  embedding_config)).with_transform(scale_config)
 
       def assert_element(element):
-        assert max(element.feature_1) == 1
+        embedding_values = element.embedding
+        assert 0 <= max(embedding_values) <= 1
+        assert 0 <= min(embedding_values) <= 1
 
       _ = (transformed_pcoll | beam.Map(assert_element))
 
@@ -186,7 +187,7 @@ def test_with_int_data_types(self):
                 write_artifact_location=self.artifact_location).with_transform(
                     embedding_config))
 
-  def test_with_artifact_location(self):  # pylint: disable=line-too-long
+  def test_with_artifact_location(self):
     """Local artifact location test"""
     secondary_artifact_location = tempfile.mkdtemp(
         prefix='_openai_secondary_test')
@@ -231,7 +232,7 @@ def assert_element(element):
       # Clean up the temporary directory
       shutil.rmtree(secondary_artifact_location)
 
-  def test_mltransform_to_ptransform_with_openai(self):  # pylint: disable=line-too-long
+  def test_mltransform_to_ptransform_with_openai(self):
     transforms = [
         OpenAITextEmbeddings(
             columns=['x'],