89 lines
2.4 KiB
Rust
89 lines
2.4 KiB
Rust
use std::{
|
|
path::PathBuf,
|
|
sync::{Arc, Mutex},
|
|
};
|
|
|
|
use async_trait::async_trait;
|
|
use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
|
|
|
|
use domain::{
|
|
errors::{DomainError, DomainResult},
|
|
smart::ports::EmbeddingGenerator,
|
|
};
|
|
|
|
pub struct FastEmbedConfig {
|
|
pub model: EmbeddingModel,
|
|
/// Directory used to cache downloaded model files.
|
|
/// Defaults to the system cache directory when `None`.
|
|
pub cache_dir: Option<PathBuf>,
|
|
pub show_download_progress: bool,
|
|
}
|
|
|
|
impl Default for FastEmbedConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
model: EmbeddingModel::AllMiniLML6V2,
|
|
cache_dir: None,
|
|
show_download_progress: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl FastEmbedConfig {
|
|
pub fn with_model(model: EmbeddingModel) -> Self {
|
|
Self {
|
|
model,
|
|
..Default::default()
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct FastEmbedGenerator {
|
|
model: Arc<Mutex<TextEmbedding>>,
|
|
}
|
|
|
|
impl FastEmbedGenerator {
|
|
/// Initialise the model. Downloads and caches model files on first call.
|
|
pub fn new(config: FastEmbedConfig) -> Result<Self, DomainError> {
|
|
let mut opts = TextInitOptions::new(config.model)
|
|
.with_show_download_progress(config.show_download_progress);
|
|
|
|
if let Some(dir) = config.cache_dir {
|
|
opts = opts.with_cache_dir(dir);
|
|
}
|
|
|
|
let model = TextEmbedding::try_new(opts)
|
|
.map_err(|e| DomainError::Infrastructure(format!("fastembed init failed: {e}")))?;
|
|
|
|
Ok(Self {
|
|
model: Arc::new(Mutex::new(model)),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl EmbeddingGenerator for FastEmbedGenerator {
|
|
async fn generate(&self, text: &str) -> DomainResult<Vec<f32>> {
|
|
let model = Arc::clone(&self.model);
|
|
let text = text.to_owned();
|
|
|
|
tokio::task::spawn_blocking(move || {
|
|
let mut guard = model
|
|
.lock()
|
|
.map_err(|_| DomainError::Infrastructure("model mutex poisoned".into()))?;
|
|
guard
|
|
.embed(vec![text.as_str()], None)
|
|
.map_err(|e| DomainError::Infrastructure(format!("embedding failed: {e}")))?
|
|
.into_iter()
|
|
.next()
|
|
.ok_or_else(|| DomainError::Infrastructure("no embedding returned".into()))
|
|
})
|
|
.await
|
|
.map_err(|e| DomainError::Infrastructure(format!("spawn_blocking panicked: {e}")))?
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "tests/lib.rs"]
|
|
mod tests;
|