Add EmbeddingGemma retrieval support#147
Conversation
flupkede
left a comment
There was a problem hiding this comment.
Thanks for this — the EmbeddingGemma addition and the query/document split look solid. Before merging I'd like two adjustments, both about blast radius on existing indexes. Could you apply them on the branch? I've left an inline suggestion for the first one and the snippets for the rest below; happy to drop more inline suggestions or share my local diff if that's easier. I'll only push to the branch myself if you'd prefer that.
1. Gate the Text:/Code: label on the model.
Right now the label switches on file extension for every model, so upgrading changes the embedding input for .md/.txt chunks even on the default minilm-l6-q. On an existing index that re-embeds only changed files incrementally, that mixes two document representations in one vector space (and isn't covered by the "changing models needs a reindex" note). Since the label is really a Gemma affordance, let's move it onto ModelType so only EmbeddingGemma emits Text: and every other model keeps Code: unconditionally — existing indexes stay byte-for-byte unchanged.
Add to impl ModelType (next to prepare_query/prepare_document):
/// Label prefix for a chunk's main content when building the embedding input.
///
/// EmbeddingGemma benefits from distinguishing prose from source code, so
/// Markdown and plain-text files are labeled `Text`. Every other model keeps
/// the historical `Code` label unconditionally, leaving their embeddings —
/// and therefore existing indexes — byte-for-byte unchanged.
pub fn content_label(&self, path: &str) -> &'static str {
match self {
Self::EmbeddingGemma300MQ4 => {
match std::path::Path::new(path)
.extension()
.and_then(|extension| extension.to_str())
{
Some("md" | "markdown" | "txt") => "Text",
_ => "Code",
}
}
_ => "Code",
}
}That requires prepare_text to know the model. It's fixed for the embedder's lifetime, so read it once per call rather than per chunk — e.g. in embed_chunks:
let model_type = self
.embedder
.lock()
.map_err(|e| anyhow::anyhow!("Embedder mutex poisoned: {}", e))?
.model_type();
// ... then: .map(|chunk| self.prepare_text(chunk, model_type))and change the signature to fn prepare_text(&self, chunk: &Chunk, model_type: ModelType). See the inline suggestion for the label line itself.
2. Warn on memory for heavier models.
EmbeddingGemma is 768-dim vs the 384-dim default — 2× the vector size, so indexing/serving uses noticeably more RAM and disk, which can OOM memory-limited hosts (containers). A one-line warning when a heavier-than-default model is selected would help:
/// Whether this model produces larger embeddings than the default model.
pub fn is_heavier_than_default(&self) -> bool {
self.dimensions() > Self::default().dimensions()
}and emit a warning in the CLI when --model resolves to such a model (both the global flag and index add --model).
Both changes are cheap to cover with model-free unit tests (content_label returns Text only for Gemma on md/txt; is_heavier_than_default true for 768/1024-dim models). Thanks!
| // Add main content | ||
| parts.push(format!("Code:\n{}", chunk.content)); | ||
| let label = match std::path::Path::new(&chunk.path) | ||
| .extension() | ||
| .and_then(|extension| extension.to_str()) | ||
| { | ||
| Some("md" | "markdown" | "txt") => "Text", | ||
| _ => "Code", | ||
| }; | ||
| parts.push(format!("{label}:\n{}", chunk.content)); |
There was a problem hiding this comment.
| // Add main content | |
| parts.push(format!("Code:\n{}", chunk.content)); | |
| let label = match std::path::Path::new(&chunk.path) | |
| .extension() | |
| .and_then(|extension| extension.to_str()) | |
| { | |
| Some("md" | "markdown" | "txt") => "Text", | |
| _ => "Code", | |
| }; | |
| parts.push(format!("{label}:\n{}", chunk.content)); | |
| // Add main content. The label is model-specific: only models that | |
| // benefit from a prose/code distinction (EmbeddingGemma) use `Text`; | |
| // all others keep `Code`, so their embeddings stay unchanged. | |
| let label = model_type.content_label(&chunk.path); | |
| parts.push(format!("{label}:\n{}", chunk.content)); |
ℹ️ This suggestion assumes ModelType::content_label exists and that model_type is passed into prepare_text (see the two snippets in the review summary) — that's why the whole change can't be a single one-click suggestion.
|
Quick heads-up: I'm on leave for ~10 days without a laptop, so I won't be able to respond until around July 22. No rush on your side at all — the requested changes above have everything you need, and I'll review as soon as I'm back. Thanks for the contribution and your patience! 🌴 |
Scope prose-aware Text labels to EmbeddingGemma Markdown and plain-text chunks. Keep the historical Code label for all existing models so incremental indexing cannot mix document representations. Warn when explicitly selecting models with larger vector dimensions.
|
Thanks, updated. Existing models keep their original I also added regression tests and prevented |
|
Thanks so much for this, @markschroedr! 🙏 The implementation is solid. Our repo integrates through Closing this one in favour of #155. Really appreciate the contribution! 🚀 |
What changed
EmbeddingGemma300MQ4FastEmbed model under the CLI aliasembeddinggemma-q4Text:only for EmbeddingGemma, preserving existing model inputs--modelvalues that don't match an existing indexWhy
EmbeddingGemma is a compact multilingual retrieval model for source code, technical documentation, and note collections. It expects distinct query and document prompts, so treating both inputs identically reduces retrieval quality.
User impact
Users can create an EmbeddingGemma-backed index with:
Existing models and indexes retain their current behavior.
Validation
cargo fmt --all -- --checkgit diff --checkcargo test --locked --lib -- --test-threads=1(583 passed,12 ignored)