Skip to main content

swiftide_integrations/fastembed/
sparse_embedding_model.rs

1use async_trait::async_trait;
2use swiftide_core::chat_completion::errors::LanguageModelError;
3use swiftide_core::{SparseEmbedding, SparseEmbeddingModel, SparseEmbeddings};
4
5use super::{EmbeddingModelType, FastEmbed};
6#[async_trait]
7impl SparseEmbeddingModel for FastEmbed {
8    #[tracing::instrument(skip_all)]
9    async fn sparse_embed(
10        &self,
11        input: Vec<String>,
12    ) -> Result<SparseEmbeddings, LanguageModelError> {
13        let mut embedding_model = self.embedding_model.lock().await;
14
15        match &mut *embedding_model {
16            EmbeddingModelType::Sparse(model) => model
17                .embed(input, self.batch_size)
18                .map_err(LanguageModelError::permanent)
19                .and_then(|embeddings| {
20                    embeddings
21                        .into_iter()
22                        .map(|embedding| {
23                            let indices = embedding
24                                .indices
25                                .iter()
26                                .map(|v| u32::try_from(*v).map_err(LanguageModelError::permanent))
27                                .collect::<Result<Vec<_>, LanguageModelError>>()?;
28
29                            Ok(SparseEmbedding {
30                                indices,
31                                values: embedding.values,
32                            })
33                        })
34                        .collect()
35                }),
36            EmbeddingModelType::Dense(_) => Err(LanguageModelError::PermanentError(
37                "Expected sparse model, got dense".into(),
38            )),
39        }
40    }
41}