postgresml · SilasMarvin · Nov 9, 2023 · Nov 8, 2023 · Nov 8, 2023 · Nov 8, 2023
diff --git a/pgml-sdks/pgml/Cargo.toml b/pgml-sdks/pgml/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pgml"
-version = "0.9.5"
+version = "0.9.6"
 edition = "2021"
 authors = ["PosgresML <team@postgresml.org>"]
 homepage = "https://postgresml.org/"

diff --git a/pgml-sdks/pgml/javascript/package.json b/pgml-sdks/pgml/javascript/package.json
@@ -1,6 +1,6 @@
 {
   "name": "pgml",
-  "version": "0.9.5",
+  "version": "0.9.6",
   "description": "Open Source Alternative for Building End-to-End Vector Search Applications without OpenAI & Pinecone",
   "keywords": [
     "postgres",

diff --git a/pgml-sdks/pgml/javascript/tests/typescript-tests/test.ts b/pgml-sdks/pgml/javascript/tests/typescript-tests/test.ts
@@ -280,6 +280,28 @@ it("can order documents", async () => {
   await collection.archive();
 });
 
+///////////////////////////////////////////////////
+// Transformer Pipeline Tests /////////////////////
+///////////////////////////////////////////////////
+
+it("can transformer pipeline", async () => {
+  const t = pgml.newTransformerPipeline("text-generation");
+  const it = await t.transform(["AI is going to"], {max_new_tokens: 5});
+  expect(it.length).toBeGreaterThan(0)
+});
+
+it("can transformer pipeline stream", async () => {
+  const t = pgml.newTransformerPipeline("text-generation");
+  const it = await t.transform_stream("AI is going to", {max_new_tokens: 5});
+  let result = await it.next();
+  let output = [];
+  while (!result.done) {
+    output.push(result.value);
+    result = await it.next();
+  }
+  expect(output.length).toBeGreaterThan(0)
+});
+
 ///////////////////////////////////////////////////
 // Test migrations ////////////////////////////////
 ///////////////////////////////////////////////////

diff --git a/pgml-sdks/pgml/pyproject.toml b/pgml-sdks/pgml/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "maturin"
 [project]
 name = "pgml"
 requires-python = ">=3.7"
-version = "0.9.5"
+version = "0.9.6"
 description = "Python SDK is designed to facilitate the development of scalable vector search applications on PostgreSQL databases."
 authors = [
   {name = "PostgresML",  email = "team@postgresml.org"},

diff --git a/pgml-sdks/pgml/python/pgml/pgml.pyi b/pgml-sdks/pgml/python/pgml/pgml.pyi
diff --git a/pgml-sdks/pgml/python/tests/test.py b/pgml-sdks/pgml/python/tests/test.py
@@ -298,6 +298,27 @@ async def test_order_documents():
     await collection.archive()
 
 
+###################################################
+## Transformer Pipeline Tests #####################
+###################################################
+
+
+@pytest.mark.asyncio
+async def test_transformer_pipeline():
+    t = pgml.TransformerPipeline("text-generation")
+    it = await t.transform(["AI is going to"], {"max_new_tokens": 5})
+    assert (len(it)) > 0
+
+@pytest.mark.asyncio
+async def test_transformer_pipeline_stream():
+    t = pgml.TransformerPipeline("text-generation")
+    it = await t.transform_stream("AI is going to", {"max_new_tokens": 5})
+    total = []
+    async for c in it:
+        total.append(c)
+    assert (len(total)) > 0
+
+
 ###################################################
 ## Migration tests ################################
 ###################################################

diff --git a/pgml-sdks/pgml/src/languages/javascript.rs b/pgml-sdks/pgml/src/languages/javascript.rs
@@ -1,8 +1,11 @@
+use futures::StreamExt;
 use neon::prelude::*;
 use rust_bridge::javascript::{FromJsType, IntoJsResult};
+use std::sync::Arc;
 
 use crate::{
     pipeline::PipelineSyncData,
+    transformer_pipeline::TransformerStream,
     types::{DateTime, Json},
 };
 
@@ -16,8 +19,9 @@ impl IntoJsResult for DateTime {
         self,
         cx: &mut C,
     ) -> JsResult<'b, Self::Output> {
-        let date = neon::types::JsDate::new(cx, self.0.assume_utc().unix_timestamp() as f64 * 1000.0)
-            .expect("Error converting to JS Date");
+        let date =
+            neon::types::JsDate::new(cx, self.0.assume_utc().unix_timestamp() as f64 * 1000.0)
+                .expect("Error converting to JS Date");
         Ok(date)
     }
 }
@@ -69,6 +73,64 @@ impl IntoJsResult for PipelineSyncData {
     }
 }
 
+#[derive(Clone)]
+struct TransformerStreamArcMutex(Arc<tokio::sync::Mutex<TransformerStream>>);
+
+impl Finalize for TransformerStreamArcMutex {}
+
+fn transform_stream_iterate_next(mut cx: FunctionContext) -> JsResult<JsPromise> {
+    let this = cx.this();
+    let s: Handle<JsBox<TransformerStreamArcMutex>> = this
+        .get(&mut cx, "s")
+        .expect("Error getting self in transformer_stream_iterate_next");
+    let ts: &TransformerStreamArcMutex = &s;
+    let ts: TransformerStreamArcMutex = ts.clone();
+
+    let channel = cx.channel();
+    let (deferred, promise) = cx.promise();
+    crate::get_or_set_runtime().spawn(async move {
+        let mut ts = ts.0.lock().await;
+        let v = ts.next().await;
+        deferred
+            .try_settle_with(&channel, move |mut cx| {
+                let o = cx.empty_object();
+                if let Some(v) = v {
+                    let v: String = v.expect("Error calling next on TransformerStream");
+                    let v = cx.string(v);
+                    let d = cx.boolean(false);
+                    o.set(&mut cx, "value", v)
+                        .expect("Error setting object value in transformer_sream_iterate_next");
+                    o.set(&mut cx, "done", d)
+                        .expect("Error setting object value in transformer_sream_iterate_next");
+                } else {
+                    let d = cx.boolean(true);
+                    o.set(&mut cx, "done", d)
+                        .expect("Error setting object value in transformer_sream_iterate_next");
+                }
+                Ok(o)
+            })
+            .expect("Error sending js");
+    });
+    Ok(promise)
+}
+
+impl IntoJsResult for TransformerStream {
+    type Output = JsObject;
+    fn into_js_result<'a, 'b, 'c: 'b, C: Context<'c>>(
+        self,
+        cx: &mut C,
+    ) -> JsResult<'b, Self::Output> {
+        let o = cx.empty_object();
+        let f: Handle<JsFunction> = JsFunction::new(cx, transform_stream_iterate_next)?;
+        o.set(cx, "next", f)?;
+        let s = cx.boxed(TransformerStreamArcMutex(Arc::new(
+            tokio::sync::Mutex::new(self),
+        )));
+        o.set(cx, "s", s)?;
+        Ok(o)
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // JS To Rust //////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/pgml-sdks/pgml/src/languages/python.rs b/pgml-sdks/pgml/src/languages/python.rs
@@ -1,65 +1,99 @@
+use futures::StreamExt;
 use pyo3::conversion::IntoPy;
 use pyo3::types::{PyDict, PyFloat, PyInt, PyList, PyString};
 use pyo3::{prelude::*, types::PyBool};
+use std::sync::Arc;
 
 use rust_bridge::python::CustomInto;
 
-use crate::{pipeline::PipelineSyncData, types::Json};
+use crate::{pipeline::PipelineSyncData, transformer_pipeline::TransformerStream, types::Json};
 
 ////////////////////////////////////////////////////////////////////////////////
 // Rust to PY //////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 
-impl ToPyObject for Json {
-    fn to_object(&self, py: Python) -> PyObject {
+impl IntoPy<PyObject> for Json {
+    fn into_py(self, py: Python) -> PyObject {
         match &self.0 {
-            serde_json::Value::Bool(x) => x.to_object(py),
+            serde_json::Value::Bool(x) => x.into_py(py),
             serde_json::Value::Number(x) => {
                 if x.is_f64() {
                     x.as_f64()
                         .expect("Error converting to f64 in impl ToPyObject for Json")
-                        .to_object(py)
+                        .into_py(py)
                 } else {
                     x.as_i64()
                         .expect("Error converting to i64 in impl ToPyObject for Json")
-                        .to_object(py)
+                        .into_py(py)
                 }
             }
-            serde_json::Value::String(x) => x.to_object(py),
+            serde_json::Value::String(x) => x.into_py(py),
             serde_json::Value::Array(x) => {
                 let list = PyList::empty(py);
                 for v in x.iter() {
-                    list.append(Json(v.clone()).to_object(py)).unwrap();
+                    list.append(Json(v.clone()).into_py(py)).unwrap();
                 }
-                list.to_object(py)
+                list.into_py(py)
             }
             serde_json::Value::Object(x) => {
                 let dict = PyDict::new(py);
                 for (k, v) in x.iter() {
-                    dict.set_item(k, Json(v.clone()).to_object(py)).unwrap();
+                    dict.set_item(k, Json(v.clone()).into_py(py)).unwrap();
                 }
-                dict.to_object(py)
+                dict.into_py(py)
             }
             serde_json::Value::Null => py.None(),
         }
     }
 }
 
-impl IntoPy<PyObject> for Json {
+impl IntoPy<PyObject> for PipelineSyncData {
     fn into_py(self, py: Python) -> PyObject {
-        self.to_object(py)
+        Json::from(self).into_py(py)
     }
 }
 
-impl ToPyObject for PipelineSyncData {
-    fn to_object(&self, py: Python) -> PyObject {
-        Json::from(self.clone()).to_object(py)
+#[pyclass]
+#[derive(Clone)]
+struct TransformerStreamPython {
+    wrapped: Arc<tokio::sync::Mutex<TransformerStream>>,
+}
+
+#[pymethods]
+impl TransformerStreamPython {
+    fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __anext__<'p>(slf: PyRefMut<'_, Self>, py: Python<'p>) -> PyResult<Option<PyObject>> {
+        let ts = slf.wrapped.clone();
+        let fut = pyo3_asyncio::tokio::future_into_py(py, async move {
+            let mut ts = ts.lock().await;
+            if let Some(o) = ts.next().await {
+                Ok(Some(Python::with_gil(|py| {
+                    o.expect("Error calling next on TransformerStream")
+                        .to_object(py)
+                })))
+            } else {
+                Err(pyo3::exceptions::PyStopAsyncIteration::new_err(
+                    "stream exhausted",
+                ))
+            }
+        })?;
+        Ok(Some(fut.into()))
     }
 }
 
-impl IntoPy<PyObject> for PipelineSyncData {
+impl IntoPy<PyObject> for TransformerStream {
     fn into_py(self, py: Python) -> PyObject {
-        self.to_object(py)
+        let f: Py<TransformerStreamPython> = Py::new(
+            py,
+            TransformerStreamPython {
+                wrapped: Arc::new(tokio::sync::Mutex::new(self)),
+            },
+        )
+        .expect("Error converting TransformerStream to TransformerStreamPython");
+        f.to_object(py)
     }
 }
 
@@ -115,6 +149,12 @@ impl FromPyObject<'_> for PipelineSyncData {
     }
 }
 
+impl FromPyObject<'_> for TransformerStream {
+    fn extract(_ob: &PyAny) -> PyResult<Self> {
+        panic!("We must implement this, but this is impossible to be reached")
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Rust to Rust //////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////