Updated some signatures added some very basic tests and cleaned some …

…stuff up
postgresml · SilasMarvin · Nov 9, 2023 · Nov 8, 2023 · Nov 8, 2023 · Nov 8, 2023
commit 231e2292e18deb867febc59b5427f39bd0ac4fdd
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/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
@@ -83,7 +83,6 @@ fn transform_stream_iterate_next(mut cx: FunctionContext) -> JsResult<JsPromise>
     let s: Handle<JsBox<TransformerStreamArcMutex>> = this
         .get(&mut cx, "s")
         .expect("Error getting self in transformer_stream_iterate_next");
-    let mut b: &JsBox<TransformerStreamArcMutex> = &s;
     let ts: &TransformerStreamArcMutex = &s;
     let ts: TransformerStreamArcMutex = ts.clone();
 

diff --git a/pgml-sdks/pgml/src/languages/python.rs b/pgml-sdks/pgml/src/languages/python.rs
@@ -65,7 +65,7 @@ impl TransformerStreamPython {
         slf
     }
 
-    fn __anext__<'p>(mut slf: PyRefMut<'_, Self>, py: Python<'p>) -> PyResult<Option<PyObject>> {
+    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;
@@ -150,7 +150,7 @@ impl FromPyObject<'_> for PipelineSyncData {
 }
 
 impl FromPyObject<'_> for TransformerStream {
-    fn extract(ob: &PyAny) -> PyResult<Self> {
+    fn extract(_ob: &PyAny) -> PyResult<Self> {
         panic!("We must implement this, but this is impossible to be reached")
     }
 }

diff --git a/pgml-sdks/pgml/src/transformer_pipeline.rs b/pgml-sdks/pgml/src/transformer_pipeline.rs
@@ -163,7 +163,7 @@ impl TransformerPipeline {
     #[instrument(skip(self))]
     pub async fn transform_stream(
         &self,
-        inputs: Vec<String>,
+        input: &str,
         args: Option<Json>,
         batch_size: Option<i32>,
     ) -> anyhow::Result<TransformerStream> {
@@ -173,10 +173,10 @@ impl TransformerPipeline {
 
         let mut transaction = pool.begin().await?;
         sqlx::query(
-            "DECLARE c CURSOR FOR SELECT pgml.transform_stream(task => $1, inputs => $2, args => $3)",
+            "DECLARE c CURSOR FOR SELECT pgml.transform_stream(task => $1, input => $2, args => $3)",
         )
         .bind(&self.task)
-        .bind(inputs)
+        .bind(input)
         .bind(&args)
         .execute(&mut *transaction)
         .await?;
@@ -246,7 +246,7 @@ mod tests {
         );
         let mut stream = t
             .transform_stream(
-                vec!["AI is going to".to_string()],
+                "AI is going to",
                 Some(
                     serde_json::json!({
                         "max_new_tokens": 10