Silas add ranking #1498
Silas add ranking #1498
Conversation
|
|
||
| if transformer not in __cache_sentence_transformer_by_name: | ||
| __cache_sentence_transformer_by_name[transformer] = create_cross_encoder( | ||
| transformer |
There was a problem hiding this comment.
If you pass kwargs through to create_cross_encoder we can specify the device, https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html?highlight=crossencoder#sentence_transformers.cross_encoder.CrossEncoder
We should do this for the SentenceTransformer constructor, too.
There was a problem hiding this comment.
We will have to create a separate argument for this, or pop specific arguments from kwargs. If just pass kwargs straight through we will get an unexpected keyword argument error.
There was a problem hiding this comment.
We should bump the version (probably to 2.9.0 since it's an API addition) and add the following migration file sql/pgml--2.8.5--2.9.0.sql:
-- src/api.rs:613
-- pgml::api::rank
CREATE FUNCTION pgml."rank"(
"transformer" TEXT, /* &str */
"query" TEXT, /* &str */
"documents" TEXT[], /* alloc::vec::Vec<&str> */
"kwargs" jsonb DEFAULT '{}' /* pgrx::datum::json::JsonB */
) RETURNS TABLE (
"corpus_id" bigint, /* i64 */
"score" double precision, /* f64 */
"text" TEXT /* core::option::Option<alloc::string::String> */
)
IMMUTABLE STRICT PARALLEL SAFE
LANGUAGE c /* Rust */
AS 'MODULE_PATHNAME', 'rank_wrapper';
Done! 745b190 |
This integration brings reranking into postgresml.