Skip to content

Add EmbeddingGemma retrieval support#147

Closed
markschroedr wants to merge 3 commits into
flupkede:masterfrom
markschroedr:agent/add-embeddinggemma-support
Closed

Add EmbeddingGemma retrieval support#147
markschroedr wants to merge 3 commits into
flupkede:masterfrom
markschroedr:agent/add-embeddinggemma-support

Conversation

@markschroedr

@markschroedr markschroedr commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

What changed

  • add the quantized EmbeddingGemma300MQ4 FastEmbed model under the CLI alias embeddinggemma-q4
  • apply EmbeddingGemma's recommended query and document prompts through separate embedding paths
  • label Markdown and plain-text chunks as Text: only for EmbeddingGemma, preserving existing model inputs
  • warn for heavier models and reject --model values that don't match an existing index
  • document model selection and the required full reindex when changing vector spaces

Why

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:

codesearch --model embeddinggemma-q4 index /path/to/project --force

Existing models and indexes retain their current behavior.

Validation

  • cargo fmt --all -- --check
  • git diff --check
  • cargo test --locked --lib -- --test-threads=1 (583 passed, 12 ignored)
  • 3 CLI regression tests
  • Clippy with warnings denied
  • local end-to-end Markdown indexing and semantic-search smoke test

@flupkede flupkede left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!

Comment thread src/embed/batch.rs Outdated
Comment on lines +180 to +188
// 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));

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// 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.

@flupkede flupkede self-assigned this Jul 11, 2026
@flupkede flupkede added the enhancement New feature or request label Jul 11, 2026
@flupkede

Copy link
Copy Markdown
Owner

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.
@markschroedr

Copy link
Copy Markdown
Contributor Author

Thanks, updated. Existing models keep their original Code: formatting; only EmbeddingGemma distinguishes text via the existing language classifier.

I also added regression tests and prevented --model from searching or syncing an incompatible index. All tests and Clippy pass.

@flupkede

Copy link
Copy Markdown
Owner

Thanks so much for this, @markschroedr! 🙏 The implementation is solid. Our repo integrates through develop (not master) via a develop-first gitflow, so I've rebased your work cleanly onto develop and opened #155 to land it there — with your authorship preserved on the commits and full credit to you. It also folds in the review feedback (prose/code labelling scoped to EmbeddingGemma only, so existing indexes are unaffected, plus a memory guard when a heavier model is selected).

Closing this one in favour of #155. Really appreciate the contribution! 🚀

@flupkede flupkede closed this Jul 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants