swiftide_integrations/fastembed/
sparse_embedding_model.rs

1use anyhow::{Context as _, Result};
2use async_trait::async_trait;
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(&self, input: Vec<String>) -> Result<SparseEmbeddings> {
10        if let EmbeddingModelType::Sparse(embedding_model) = &*self.embedding_model {
11            embedding_model
12                .embed(input, self.batch_size)
13                .and_then(|embeddings| {
14                    embeddings
15                        .into_iter()
16                        .map(|embedding| {
17                            Ok(SparseEmbedding {
18                                indices: embedding
19                                    .indices
20                                    .iter()
21                                    .map(|v| {
22                                        u32::try_from(*v).context(
23                                            "Could not convert sparse vector from u32 to usize",
24                                        )
25                                    })
26                                    .collect::<Result<Vec<_>>>()?,
27                                values: embedding.values,
28                            })
29                        })
30                        .collect()
31                })
32        } else {
33            Err(anyhow::anyhow!("Expected dense model, got sparse"))
34        }
35    }
36}