semtree_embed/
fastembed.rs1use 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}