Skip to content

Added embed2 which returns a table structure #1186

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.

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Added embed2 which returns a table structure #1186

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
503dc71
Added pyarrow==11.0.0 to requirements to solve issue where postgres w…
ns1000 Nov 23, 2023
52e5b8d
Merge branch 'master' of https://github.com/postgresml/postgresml
ns1000 Nov 25, 2023
3462868
Added embed_batch2, which will batch process multiple strings and the…
ns1000 Nov 25, 2023
c5c8640
Merge branch 'master' of https://github.com/postgresml/postgresml
ns1000 Nov 25, 2023
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
20 changes: 20 additions & 0 deletions pgml-extension/src/api.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,26 @@ pub fn embed_batch(
}
}

#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
#[pg_extern(immutable, parallel_safe, name = "embed2")]
pub fn embed_batch2<'a>(
Copy link
Contributor

@montanalow montanalow Nov 25, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My rough thoughts, without running the code on some examples.

I think we should name this SQL embed_3 and Rust embed_batch_3 with the goal of establishing this as the 3.0 embed API, as well as a pattern for releasing 3.0 APIs early as we developing them in an alpha state (with potentially breaking changes, where we completely drop them in 3.1 in favor of the newly established default behavior).

Your example convinces me that batch APIs should return a table, but I think that table's rows should be JSONB with {id, embedding} keys (at least), unless there is a significant performance implication on that front. My thinking is that embedding models are getting more complicated and now some take JSON rather than TEXT for inputs including a prompt. It would be nice to have an optional id in the input JSON, and if it's not present, then just return the entire input JSON as the id, which acts just like your TEXT as the key.

Final thought is that kwargs is JSONB currently, which works well with the underlying Python dependencies, but I'd like to structure it as much as possible for final 3.0. We should find a way to flag this obviously as an alpha API, that will be broken and eventually dropped when a final version is available.

transformer: &str,
inputs: Vec<&'a str>,
kwargs: default!(JsonB, "'{}'"),
) -> TableIterator<'a, (name!(text, String), name!(embedding, Vec<f32>))> {
let rows = match crate::bindings::transformers::embed(transformer, inputs.clone(), &kwargs.0) {
Ok(rows) => rows,
Err(e) => {
error!("{e}");
}
};
TableIterator::new(
inputs.into_iter().zip(rows.into_iter()).map(|(text, embedding)| {
(text.to_string(), embedding)
}),
)
}

/// Clears the GPU cache.
///
/// # Arguments
Expand Down