From 275bfb58085bd83cb36ea36e08287c4f5667be62 Mon Sep 17 00:00:00 2001 From: SilasMarvin <19626586+SilasMarvin@users.noreply.github.com> Date: Thu, 14 Mar 2024 10:35:37 -0700 Subject: [PATCH] Added boosting for vector search --- pgml-sdks/pgml/src/lib.rs | 8 +++++--- pgml-sdks/pgml/src/vector_search_query_builder.rs | 7 +++++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/pgml-sdks/pgml/src/lib.rs b/pgml-sdks/pgml/src/lib.rs index 34b02ce53..8f6039b52 100644 --- a/pgml-sdks/pgml/src/lib.rs +++ b/pgml-sdks/pgml/src/lib.rs @@ -1052,7 +1052,7 @@ mod tests { #[tokio::test] async fn can_vector_search_with_local_embeddings() -> anyhow::Result<()> { internal_init_logger(None, None).ok(); - let collection_name = "test r_c_cvswle_9"; + let collection_name = "test r_c_cvswle_13"; let mut collection = Collection::new(collection_name, None)?; let documents = generate_dummy_documents(10); collection.upsert_documents(documents.clone(), None).await?; @@ -1095,10 +1095,12 @@ mod tests { "parameters": { "instruction": "Represent the Wikipedia document for retrieval" }, - "full_text_filter": "test" + "full_text_filter": "test", + "boost": 1.2 }, "body": { - "query": "Test document: 2" + "query": "Test document: 2", + "boost": 1.0 }, }, "filter": { diff --git a/pgml-sdks/pgml/src/vector_search_query_builder.rs b/pgml-sdks/pgml/src/vector_search_query_builder.rs index df4f54e79..6c0381b19 100644 --- a/pgml-sdks/pgml/src/vector_search_query_builder.rs +++ b/pgml-sdks/pgml/src/vector_search_query_builder.rs @@ -25,6 +25,7 @@ struct ValidField { query: String, parameters: Option, full_text_filter: Option, + boost: Option, } #[derive(Debug, Deserialize)] @@ -98,6 +99,8 @@ pub async fn build_vector_search_query( let mut query = Query::select(); + let boost = vf.boost.unwrap_or(1.); + match model_runtime { ModelRuntime::Python => { // Build the embedding CTE @@ -118,7 +121,7 @@ pub async fn build_vector_search_query( query .expr(Expr::cust(format!( - r#"1 - (embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector) AS score"# + r#"(1 - (embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector)) * {boost} AS score"# ))) .order_by_expr(Expr::cust(format!( r#"embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector"# @@ -152,7 +155,7 @@ pub async fn build_vector_search_query( // Build the score CTE query .expr(Expr::cust_with_values( - r#"1 - (embeddings.embedding <=> $1::vector) AS score"#, + r#"(1 - (embeddings.embedding <=> $1::vector)) {boost} AS score"#, [embedding.clone()], )) .order_by_expr( 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