diff --git a/pgml-sdks/pgml/src/lib.rs b/pgml-sdks/pgml/src/lib.rs index d146a38fd..db69f52e7 100644 --- 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::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; 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 index 0a4f62d24..49197ecf1 100644 --- 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, pub(crate) database_data: Option, - // This database_url is specifically used only for the model when calling transform and other - // one-off methods - database_url: Option, } 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, - source: Option, - parameters: Option, - database_url: Option, - ) -> Self { + pub fn new(name: Option, source: Option, parameters: Option) -> 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, - args: Option, - ) -> anyhow::Result { - 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::(0); - Ok(Json(results)) - } } impl From for Model { @@ -228,7 +193,6 @@ impl From for Model { id: x.model_id, created_at: x.model_created_at, }), - database_url: None, } } } @@ -244,36 +208,6 @@ impl From 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 new file mode 100644 index 000000000..2c713ed81 --- /dev/null +++ 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, +} + +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, + args: Option, + database_url: Option, + ) -> 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, args: Option) -> anyhow::Result { + 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::(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(()) + } +} pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy