Question: Batched encoding in SentenceEmbeddings

[pratik-mahamuni]

Hello,
Thank you for reading my question. Great library.

I was wondering if there is a feature for batched encoding when using Sentence Embedding.

Depending on the input, I could be encoding ~10 to ~15,000 sentences. My computer runs out of memory (32 GB) when I was trying to transform a case with ~9000 sentences. The code is as follows:

let blocking_task = tokio::task::spawn_blocking(move || {
    let model = SentenceEmbeddingsBuilder::remote(SentenceEmbeddingsModelType::AllMiniLmL12V2)
        .create_model()
        .expect("Could not create sentence embeddings model");

    let query = [query];
    let sentences = sentences;
    let sentences_encoding = model.encode(&sentences).expect("Could not embed sentences"); // <-- Variable length
    let query_encoding = model.encode(&query).expect("Could not embed query");

    let mut map: HashMap<String, Vec<Vec<f32>>> = HashMap::new();
    map.insert("sentences".to_string(), sentences_encoding);
    map.insert("query".to_string(), query_encoding);

    map
});

let embeddings_map = blocking_task.await.unwrap();

[guillaume-be]

Hello @pratik-mahamuni ,

You are correct that the pipelines generally don't handle the batching of inputs. Given that each application and target infrastructure is different it is currently left to the application developer to implement this logic.

I invite you to check the https://crates.io/crates/batched-fn crate and the example provided at https://github.com/epwalsh/rust-dl-webserver/blob/master/src/main.rs on how to create a webserver leveraging Tokio and a batching mechanism to protect the server from running out of memory.

[pratik-mahamuni]

This is very good. Thank you for providing me the information and an example (love that).