run text encoder functions in background thread for responsivness
This commit is contained in:
parent
a62b0aacf5
commit
7181bd5ffb
1 changed files with 22 additions and 4 deletions
26
src/query.rs
26
src/query.rs
|
@ -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)
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue