Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Added TransformerPipeline#1128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Merged
SilasMarvin merged 1 commit intomasterfromsilas-transformer-pipeline
Oct 27, 2023
Merged
Show file tree
Hide file tree
Changes fromall commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletionspgml-sdks/pgml/src/lib.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@ mod query_builder;
mod query_runner;
mod remote_embeddings;
mod splitter;
mod transformer_pipeline;
pub mod types;
mod utils;

Expand All@@ -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";
Expand DownExpand Up@@ -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(())
}

Expand DownExpand Up@@ -193,6 +196,10 @@ fn main(mut cx: neon::context::ModuleContext) -> neon::result::NeonResult<()> {
cx.export_function("newModel", model::Model#"4399ecc1c68b79a9cae003f3687afeae26862b43f7ae76c535d601693dae124e"> cx.export_function("newSplitter", splitter::Splitter#"4399ecc1c68b79a9cae003f3687afeae26862b43f7ae76c535d601693dae124e"> cx.export_function("newBuiltins", builtins::Builtins#"4399ecc1c68b79a9cae003f3687afeae26862b43f7ae76c535d601693dae124e"> cx.export_function(
"newTransformerPipeline",
transformer_pipeline::TransformerPipeline#"4399ecc1c68b79a9cae003f3687afeae26862b43f7ae76c535d601693dae124e"> )?;
cx.export_function("newPipeline", pipeline::Pipeline#"4399ecc1c68b79a9cae003f3687afeae26862b43f7ae76c535d601693dae124e"> Ok(())
}
Expand DownExpand Up@@ -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(
Expand DownExpand Up@@ -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(
Expand DownExpand Up@@ -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(
Expand DownExpand Up@@ -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(
Expand Down
70 changes: 2 additions & 68 deletionspgml-sdks/pgml/src/model.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -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::{
Expand DownExpand Up@@ -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)
}
}

Expand All@@ -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());
Expand All@@ -105,7 +95,6 @@ impl Model {
parameters,
project_info: None,
database_data: None,
database_url,
}
}

Expand DownExpand Up@@ -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 {
Expand All@@ -228,7 +193,6 @@ impl From<models::PipelineWithModelAndSplitter> for Model {
id: x.model_id,
created_at: x.model_created_at,
}),
database_url: None,
}
}
}
Expand All@@ -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(())
}
}
97 changes: 97 additions & 0 deletionspgml-sdks/pgml/src/transformer_pipeline.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -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(())
}
}
[8]ページ先頭
©2009-2025 Movatter.jp