Refactor OpenAIEmbeddings

python/com/johnsnowlabs/ml/ai/__init__.py

@@ -0,0 +1,10 @@
+import sys
+
+if sys.version_info[0] == 2:
+    raise ImportError(
+        "Spark NLP for Python 2.x is deprecated since version >= 4.0. "
+        "Please use an older versions to use it with this Python version."
+    )
+else:
+    import sparknlp
+    sys.modules['com.johnsnowlabs.ml.ai'] = sparknlp

python/sparknlp/__init__.py

@@ -58,6 +58,7 @@
 sys.modules['com.johnsnowlabs.nlp.annotators.coref'] = annotator
 sys.modules['com.johnsnowlabs.nlp.annotators.cv'] = annotator
 sys.modules['com.johnsnowlabs.nlp.annotators.audio'] = annotator
+sys.modules['com.johnsnowlabs.ml.ai'] = annotator
 
 annotators = annotator
 embeddings = annotator

python/test/annotator/embeddings/open_ai_embeddings_test.py

@@ -0,0 +1,47 @@
+import unittest
+from sparknlp.annotator import *
+from sparknlp.base import *
+from pyspark.sql import DataFrame
+from pyspark.sql import SparkSession
+
+class OpenAIEmbeddingsTestCase(unittest.TestCase):
+# Set your OpenAI API key to run unit test...
+    def setUp(self):
+        self.spark = SparkSession.builder \
+            .appName("Tests") \
+            .master("local[*]") \
+            .config("spark.driver.memory","8G") \
+            .config("spark.driver.maxResultSize", "2G") \
+            .config("spark.jars", "lib/sparknlp.jar") \
+            .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \
+            .config("spark.kryoserializer.buffer.max", "1000m") \
+            .config("spark.jsl.settings.openai.api.key","") \
+            .getOrCreate()
+
+    def test_openai_embeddings(self):
+
+        documentAssembler = DocumentAssembler() \
+                 .setInputCol("text") \
+                 .setOutputCol("document")
+        openai_embeddings = OpenAIEmbeddings() \
+                 .setInputCols("document") \
+                 .setOutputCol("embeddings") \
+                 .setModel("text-embedding-ada-002")
+
+        import tempfile
+        openai_embeddings.write().overwrite().save("file:///" + tempfile.gettempdir() + "/openai_embeddings")
+        loaded = OpenAIEmbeddings.load("file:///" + tempfile.gettempdir() + "/openai_embeddings")
+
+        pipeline = Pipeline().setStages([
+             documentAssembler,
+            loaded
+         ])
+
+        sample_text = [["The food was delicious and the waiter..."]]
+        sample_df = self.spark.createDataFrame(sample_text).toDF("text")
+        pipeline.fit(sample_df).transform(sample_df).select("embeddings").show(truncate=False)
+
+
+
+if __name__ == '__main__':
+    unittest.main()

src/test/scala/com/johnsnowlabs/ml/ai/OpenAIEmbeddingsTest.scala

@@ -1,13 +1,28 @@
 package com.johnsnowlabs.ml.ai
 
-import com.johnsnowlabs.nlp.annotators.SparkSessionTest
 import com.johnsnowlabs.tags.SlowTest
 import org.apache.spark.ml.Pipeline
+import org.apache.spark.sql.SparkSession
 import org.scalatest.flatspec.AnyFlatSpec
 
-class OpenAIEmbeddingsTest extends AnyFlatSpec with SparkSessionTest {
+class OpenAIEmbeddingsTest extends AnyFlatSpec {
+
+  private val spark = SparkSession
+    .builder()
+    .appName("test")
+    .master("local[*]")
+    .config("spark.driver.memory", "16G")
+    .config("spark.driver.maxResultSize", "0")
+    .config("spark.kryoserializer.buffer.max", "2000M")
+    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
+    .config("spark.jsl.settings.openai.api.key",
+      "" // Set your OpenAI API key here...
+    )
+    .getOrCreate()
 
   import spark.implicits._
+  private val documentAssembler =
+    new com.johnsnowlabs.nlp.DocumentAssembler().setInputCol("text").setOutputCol("document")
 
   "OpenAIEmbeddings" should "generate a completion for prompts" taggedAs SlowTest in {
     // Set OPENAI_API_KEY env variable to make this work
@@ -25,4 +40,26 @@ class OpenAIEmbeddingsTest extends AnyFlatSpec with SparkSessionTest {
     completionDF.select("embeddings").show(false)
   }
 
+  "OpenAIEmbeddings" should "work with escape chars" taggedAs SlowTest in {
+    val data = Seq(
+      (1, "Hello \"World\""),
+      (2, "Hello \n World"),
+      (3, "Hello \t World"),
+      (4, "Hello \r World"),
+      (5, "Hello \b World"),
+      (6, "Hello \f World"),
+      (7, "Hello \\ World"))
+    val columns = Seq("id", "text")
+    val testDF = spark.createDataFrame(data).toDF(columns: _*)
+
+    val openAIEmbeddings = new OpenAIEmbeddings()
+      .setInputCols("document")
+      .setOutputCol("embeddings")
+      .setModel("text-embedding-ada-002")
+
+    val pipeline = new Pipeline().setStages(Array(documentAssembler, openAIEmbeddings))
+    val resultDF = pipeline.fit(testDF).transform(testDF)
+    resultDF.select("embeddings").show(false)
+  }
+
 }