embed_anything 0.6.7

Embed anything at lightning speed
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

use anyhow::Error as E;
use model2vec_rs;

use crate::embeddings::embed::EmbeddingResult;

pub struct Model2VecEmbedder {
    pub model: model2vec_rs::model::StaticModel,
}

impl Model2VecEmbedder {
    pub fn new(model_id: &str, token: Option<&str>, path_in_repo: Option<&str>) -> Result<Self, E> {
        let model =
            model2vec_rs::model::StaticModel::from_pretrained(model_id, token, None, path_in_repo)?;
        Ok(Self { model })
    }

    pub fn embed(
        &self,
        text_batch: &[&str],
        _batch_size: Option<usize>,
    ) -> Result<Vec<EmbeddingResult>, E> {
        let embeddings = self.model.encode(
            text_batch
                .iter()
                .map(|s| s.to_string())
                .collect::<Vec<String>>()
                .as_slice(),
        );
        let embeddings = embeddings
            .iter()
            .map(|e| EmbeddingResult::DenseVector(e.clone()))
            .collect();
        Ok(embeddings)
    }
}