postgresml · SilasMarvin · Oct 27, 2023 · Oct 27, 2023
diff --git a/pgml-sdks/pgml/src/lib.rs b/pgml-sdks/pgml/src/lib.rs
@@ -26,6 +26,7 @@ mod query_builder;
 mod query_runner;
 mod remote_embeddings;
 mod splitter;
+mod transformer_pipeline;
 pub mod types;
 mod utils;
 
@@ -35,6 +36,7 @@ pub use collection::Collection;
 pub use model::Model;
 pub use pipeline::Pipeline;
 pub use splitter::Splitter;
+pub use transformer_pipeline::TransformerPipeline;
 
 // This is use when inserting collections to set the sdk_version used during creation
 static SDK_VERSION: &str = "0.9.2";
@@ -149,6 +151,7 @@ fn pgml(_py: pyo3::Python, m: &pyo3::types::PyModule) -> pyo3::PyResult<()> {
     m.add_class::<model::ModelPython>()?;
     m.add_class::<splitter::SplitterPython>()?;
     m.add_class::<builtins::BuiltinsPython>()?;
+    m.add_class::<transformer_pipeline::TransformerPipelinePython>()?;
     Ok(())
 }
 
@@ -193,6 +196,10 @@ fn main(mut cx: neon::context::ModuleContext) -> neon::result::NeonResult<()> {
     cx.export_function("newModel", model::ModelJavascript::new)?;
     cx.export_function("newSplitter", splitter::SplitterJavascript::new)?;
     cx.export_function("newBuiltins", builtins::BuiltinsJavascript::new)?;
+    cx.export_function(
+        "newTransformerPipeline",
+        transformer_pipeline::TransformerPipelineJavascript::new,
+    )?;
     cx.export_function("newPipeline", pipeline::PipelineJavascript::new)?;
     Ok(())
 }
@@ -448,7 +455,6 @@ mod tests {
             Some("text-embedding-ada-002".to_string()),
             Some("openai".to_string()),
             None,
-            None,
         );
         let splitter = Splitter::default();
         let mut pipeline = Pipeline::new(
@@ -527,7 +533,6 @@ mod tests {
             Some("hkunlp/instructor-base".to_string()),
             Some("python".to_string()),
             Some(json!({"instruction": "Represent the Wikipedia document for retrieval: "}).into()),
-            None,
         );
         let splitter = Splitter::default();
         let mut pipeline = Pipeline::new(
@@ -579,7 +584,6 @@ mod tests {
             Some("text-embedding-ada-002".to_string()),
             Some("openai".to_string()),
             None,
-            None,
         );
         let splitter = Splitter::default();
         let mut pipeline = Pipeline::new(
@@ -660,7 +664,6 @@ mod tests {
             Some("text-embedding-ada-002".to_string()),
             Some("openai".to_string()),
             None,
-            None,
         );
         let splitter = Splitter::default();
         let mut pipeline = Pipeline::new(

diff --git a/pgml-sdks/pgml/src/model.rs b/pgml-sdks/pgml/src/model.rs
@@ -1,8 +1,6 @@
 use anyhow::Context;
 use rust_bridge::{alias, alias_methods};
-use serde_json::json;
 use sqlx::postgres::PgPool;
-use sqlx::Row;
 use tracing::instrument;
 
 use crate::{
@@ -61,14 +59,11 @@ pub struct Model {
     pub parameters: Json,
     project_info: Option<ProjectInfo>,
     pub(crate) database_data: Option<ModelDatabaseData>,
-    // This database_url is specifically used only for the model when calling transform and other
-    // one-off methods
-    database_url: Option<String>,
 }
 
 impl Default for Model {
     fn default() -> Self {
-        Self::new(None, None, None, None)
+        Self::new(None, None, None)
     }
 }
 
@@ -88,12 +83,7 @@ impl Model {
     /// use pgml::Model;
     /// let model = Model::new(Some("intfloat/e5-small".to_string()), None, None, None);
     /// ```
-    pub fn new(
-        name: Option<String>,
-        source: Option<String>,
-        parameters: Option<Json>,
-        database_url: Option<String>,
-    ) -> Self {
+    pub fn new(name: Option<String>, source: Option<String>, parameters: Option<Json>) -> Self {
         let name = name.unwrap_or("intfloat/e5-small".to_string());
         let parameters = parameters.unwrap_or(Json(serde_json::json!({})));
         let source = source.unwrap_or("pgml".to_string());
@@ -105,7 +95,6 @@ impl Model {
             parameters,
             project_info: None,
             database_data: None,
-            database_url,
         }
     }
 
@@ -191,30 +180,6 @@ impl Model {
             .database_url;
         get_or_initialize_pool(database_url).await
     }
-
-    pub async fn transform(
-        &self,
-        task: &str,
-        inputs: Vec<String>,
-        args: Option<Json>,
-    ) -> anyhow::Result<Json> {
-        let pool = get_or_initialize_pool(&self.database_url).await?;
-        let task = json!({
-            "task": task,
-            "model": self.name,
-            "trust_remote_code": true
-        });
-        let args = args.unwrap_or_default();
-        let query = sqlx::query("SELECT pgml.transform(task => $1, inputs => $2, args => $3)");
-        let results = query
-            .bind(task)
-            .bind(inputs)
-            .bind(&args)
-            .fetch_all(&pool)
-            .await?;
-        let results = results.get(0).unwrap().get::<serde_json::Value, _>(0);
-        Ok(Json(results))
-    }
 }
 
 impl From<models::PipelineWithModelAndSplitter> for Model {
@@ -228,7 +193,6 @@ impl From<models::PipelineWithModelAndSplitter> for Model {
                 id: x.model_id,
                 created_at: x.model_created_at,
             }),
-            database_url: None,
         }
     }
 }
@@ -244,36 +208,6 @@ impl From<models::Model> for Model {
                 id: model.id,
                 created_at: model.created_at,
             }),
-            database_url: None,
         }
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::internal_init_logger;
-
-    #[sqlx::test]
-    async fn model_can_transform() -> anyhow::Result<()> {
-        internal_init_logger(None, None).ok();
-        let model = Model::new(
-            Some("Helsinki-NLP/opus-mt-en-fr".to_string()),
-            Some("pgml".to_string()),
-            None,
-            None,
-        );
-        let results = model
-            .transform(
-                "translation",
-                vec![
-                    "How are you doing today?".to_string(),
-                    "What is a good song?".to_string(),
-                ],
-                None,
-            )
-            .await?;
-        assert!(results.as_array().is_some());
-        Ok(())
-    }
-}
diff --git a/pgml-sdks/pgml/src/transformer_pipeline.rs b/pgml-sdks/pgml/src/transformer_pipeline.rs
@@ -0,0 +1,97 @@
+use rust_bridge::{alias, alias_methods};
+use sqlx::Row;
+use tracing::instrument;
+
+/// Provides access to builtin database methods
+#[derive(alias, Debug, Clone)]
+pub struct TransformerPipeline {
+    task: Json,
+    database_url: Option<String>,
+}
+
+use crate::{get_or_initialize_pool, types::Json};
+
+#[cfg(feature = "python")]
+use crate::types::JsonPython;
+
+#[alias_methods(new, transform)]
+impl TransformerPipeline {
+    pub fn new(
+        task: &str,
+        model: Option<String>,
+        args: Option<Json>,
+        database_url: Option<String>,
+    ) -> Self {
+        let mut args = args.unwrap_or_default();
+        let a = args.as_object_mut().expect("args must be an object");
+        a.insert("task".to_string(), task.to_string().into());
+        if let Some(m) = model {
+            a.insert("model".to_string(), m.into());
+        }
+
+        Self {
+            task: args,
+            database_url,
+        }
+    }
+
+    #[instrument(skip(self))]
+    pub async fn transform(&self, inputs: Vec<String>, args: Option<Json>) -> anyhow::Result<Json> {
+        let pool = get_or_initialize_pool(&self.database_url).await?;
+        let args = args.unwrap_or_default();
+
+        let results = sqlx::query("SELECT pgml.transform(task => $1, inputs => $2, args => $3)")
+            .bind(&self.task)
+            .bind(inputs)
+            .bind(&args)
+            .fetch_all(&pool)
+            .await?;
+        let results = results.get(0).unwrap().get::<serde_json::Value, _>(0);
+        Ok(Json(results))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::internal_init_logger;
+
+    #[sqlx::test]
+    async fn transformer_pipeline_can_transform() -> anyhow::Result<()> {
+        internal_init_logger(None, None).ok();
+        let t = TransformerPipeline::new(
+            "translation_en_to_fr",
+            Some("t5-base".to_string()),
+            None,
+            None,
+        );
+        let results = t
+            .transform(
+                vec![
+                    "How are you doing today?".to_string(),
+                    "What is a good song?".to_string(),
+                ],
+                None,
+            )
+            .await?;
+        assert!(results.as_array().is_some());
+        Ok(())
+    }
+
+    #[sqlx::test]
+    async fn transformer_pipeline_can_transform_with_default_model() -> anyhow::Result<()> {
+        internal_init_logger(None, None).ok();
+        let t = TransformerPipeline::new("translation_en_to_fr", None, None, None);
+        let results = t
+            .transform(
+                vec![
+                    "How are you doing today?".to_string(),
+                    "What is a good song?".to_string(),
+                ],
+                None,
+            )
+            .await?;
+        assert!(results.as_array().is_some());
+        Ok(())
+    }
+}