Adding embed_array for getting the embeddings of multiple strings (#686)

jsaied99 · web-flow · commit aad7581d92ec · 2023-06-05T16:38:10.000-07:00
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -563,9 +563,21 @@ fn load_dataset(
     TableIterator::new(vec![(name, rows)].into_iter())
 }
 
-#[pg_extern(immutable, parallel_safe)]
+#[pg_extern(immutable, parallel_safe, name = "embed")]
 pub fn embed(transformer: &str, text: &str, kwargs: default!(JsonB, "'{}'")) -> Vec<f32> {
-    crate::bindings::transformers::embed(transformer, text, &kwargs.0)
+    embed_batch(transformer, Vec::from([text]), kwargs)
+        .first()
+        .unwrap()
+        .to_vec()
+}
+
+#[pg_extern(immutable, parallel_safe, name = "embed")]
+pub fn embed_batch(
+    transformer: &str,
+    inputs: Vec<&str>,
+    kwargs: default!(JsonB, "'{}'"),
+) -> Vec<Vec<f32>> {
+    crate::bindings::transformers::embed(transformer, &inputs, &kwargs.0)
 }
 
 #[pg_extern(immutable, parallel_safe)]
diff --git a/pgml-extension/src/bindings/transformers.py b/pgml-extension/src/bindings/transformers.py
@@ -110,25 +110,30 @@ def transform(task, args, inputs):
     return json.dumps(pipe(inputs, **args), cls=NumpyJSONEncoder)
 
 
-def embed(transformer, text, kwargs):
+def embed(transformer, inputs, kwargs):
+    
+    inputs = json.loads(inputs)
     kwargs = json.loads(kwargs)
     ensure_device(kwargs)
     instructor = transformer.startswith("hkunlp/instructor")
+    
     if instructor:
         klass = INSTRUCTOR
-        text = [[kwargs.pop("instruction"), text]]
+        
+        texts_with_instructions = []
+        instruction = kwargs.pop("instruction")
+        for text in inputs:
+            texts_with_instructions.append([instruction, text])
+            
+        inputs = texts_with_instructions
     else:
         klass = SentenceTransformer
 
     if transformer not in __cache_sentence_transformer_by_name:
         __cache_sentence_transformer_by_name[transformer] = klass(transformer)
     model = __cache_sentence_transformer_by_name[transformer]
 
-    result = model.encode(text, **kwargs)
-    if instructor:
-        result = result[0]
-
-    return result
+    return model.encode(inputs, **kwargs)
 
 
 def load_dataset(name, subset, limit: None, kwargs: "{}"):
diff --git a/pgml-extension/src/bindings/transformers.rs b/pgml-extension/src/bindings/transformers.rs
@@ -35,14 +35,13 @@ pub fn transform(
     let results = Python::with_gil(|py| -> String {
         let transform: Py<PyAny> = PY_MODULE.getattr(py, "transform").unwrap().into();
 
-        let result = transform
-            .call1(
+        let result = transform.call1(
+            py,
+            PyTuple::new(
                 py,
-                PyTuple::new(
-                    py,
-                    &[task.into_py(py), args.into_py(py), inputs.into_py(py)],
-                ),
-            );
+                &[task.into_py(py), args.into_py(py), inputs.into_py(py)],
+            ),
+        );
 
         let result = match result {
             Err(e) => {
@@ -57,11 +56,12 @@ pub fn transform(
     serde_json::from_str(&results).unwrap()
 }
 
-pub fn embed(transformer: &str, text: &str, kwargs: &serde_json::Value) -> Vec<f32> {
+pub fn embed(transformer: &str, inputs: Vec<&str>, kwargs: &serde_json::Value) -> Vec<Vec<f32>> {
     crate::bindings::venv::activate();
 
     let kwargs = serde_json::to_string(kwargs).unwrap();
-    Python::with_gil(|py| -> Vec<f32> {
+    let inputs = serde_json::to_string(&inputs).unwrap();
+    Python::with_gil(|py| -> Vec<Vec<f32>> {
         let embed: Py<PyAny> = PY_MODULE.getattr(py, "embed").unwrap().into();
         embed
             .call1(
@@ -70,7 +70,7 @@ pub fn embed(transformer: &str, text: &str, kwargs: &serde_json::Value) -> Vec<f
                     py,
                     &[
                         transformer.to_string().into_py(py),
-                        text.to_string().into_py(py),
+                        inputs.into_py(py),
                         kwargs.into_py(py),
                     ],
                 ),
