Skip to main content

semtree_embed/
fastembed.rs

1use async_trait::async_trait;
2use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
3
4use crate::{EmbedError, Embedder, Embedding};
5
6pub struct FastEmbedder {
7    model: TextEmbedding,
8}
9
10impl FastEmbedder {
11    pub fn new() -> Result<Self, EmbedError> {
12        Self::with_model(EmbeddingModel::AllMiniLML6V2)
13    }
14
15    pub fn with_model(model: EmbeddingModel) -> Result<Self, EmbedError> {
16        let te = TextEmbedding::try_new(InitOptions::new(model))
17            .map_err(|e| EmbedError::ModelLoad(e.to_string()))?;
18        Ok(Self { model: te })
19    }
20}
21
22#[async_trait]
23impl Embedder for FastEmbedder {
24    async fn embed(&self, texts: &[&str]) -> Result<Vec<Embedding>, EmbedError> {
25        let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
26        self.model
27            .embed(texts, None)
28            .map_err(|e| EmbedError::EmbedFailed(e.to_string()))
29    }
30}