run text encoder functions in background thread for responsivness

This commit is contained in:
Sebastian Hugentobler 2025-07-01 21:56:28 +02:00
parent a62b0aacf5
commit 7181bd5ffb
Signed by: shu
SSH key fingerprint: SHA256:ppcx6MlixdNZd5EUM1nkHOKoyQYoJwzuQKXM6J/t66M

View file

@ -1,6 +1,7 @@
//! Query processing and document retrieval.
use snafu::{ResultExt, Snafu};
use tokio::task::JoinError;
use crate::{
storage::{
@ -18,12 +19,16 @@ pub enum AskError {
Encode { source: tokenize::EncodeError },
#[snafu(display("Failed to embed query."))]
Embed { source: text_encoder::EmbedError },
#[snafu(display("Embedding task failed to execute."))]
EmbedJoin { source: JoinError },
#[snafu(display("Failed to retrieve similar documents."))]
Query {
source: storage::queries::QueryError,
},
#[snafu(display("Failed to rerank documents."))]
Rerank { source: text_encoder::RerankError },
#[snafu(display("Reranking task failed to execute."))]
RerankJoin { source: JoinError },
}
/// Process a user query and return ranked document matches.
@ -37,15 +42,28 @@ pub async fn ask(
limit: usize,
) -> Result<Vec<DocumentMatch>, AskError> {
let encodings = tokenizer.encode(query, chunk_size).context(EncodeSnafu)?;
let embeddings = embedder.embed(encodings[0].clone()).context(EmbedSnafu)?;
let embedder = embedder.clone();
let encoding = encodings[0].clone();
let embeddings = tokio::task::spawn_blocking(move || embedder.embed(encoding))
.await
.context(EmbedJoinSnafu)?
.context(EmbedSnafu)?;
let documents = db
.query(embeddings, (limit * 10) as i32)
.await
.context(QuerySnafu)?;
let reranked_docs = reranker
.rerank(query, documents, tokenizer, limit)
.context(RerankSnafu)?;
let reranker = reranker.clone();
let tokenizer = tokenizer.clone();
let query = query.to_string();
let reranked_docs = tokio::task::spawn_blocking(move || {
reranker.rerank(&query, documents, &tokenizer, limit)
})
.await
.context(RerankJoinSnafu)?
.context(RerankSnafu)?;
Ok(reranked_docs)
}