embeddinggemma-code-search
google/embeddinggemma-300m fine-tuned for code search over a real agent codebase, using graded retrieval traces from coding-agent sessions.
This is the "custom retrieval model from agent sessions" arm of a continual-learning course: instead of hand-labeled pairs, training signal comes from graded traces — for each natural-language query about the codebase, candidate code chunks carry relevance grades (0–3) derived from what the agent actually needed. The encoder is trained so that its softmax similarity distribution over the candidates matches the grade distribution (a listwise KL-divergence loss, in the spirit of Cursor's "align the embedding space to what sessions proved relevant" approach).
Training setup
- Base model:
google/embeddinggemma-300m(sentence-transformers, Matryoshka-truncatable embeddings) - Corpus: 240 code chunks extracted from a real Python agent-harness source tree (signatures, docstrings, AST body summaries, import/call/co-edit context)
- Traces: 800 graded retrieval traces (573 train / 227 held-out), 40 candidates per trace, indirect-intent natural-language queries
- Loss: listwise KL divergence between softmax(similarities) and the normalized grade distribution
- Schedule: deliberately gentle — 1 epoch, lr 5e-6, full-parameter. (Aggressive schedules overwrite the pretrained space and hurt held-out recall.)
Results (held-out queries, recall@5)
| Embedding dim (MRL) | Base | Fine-tuned |
|---|---|---|
| 768 | 0.753 | 0.793 |
| 256 | 0.643 | 0.753 |
| 128 | 0.568 | 0.722 |
The biggest lift is at truncated Matryoshka dimensions — the fine-tuned 128-d embeddings match the base model's 256-d quality, which is what you want for cheap, low-latency code search indexes.
Usage
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("jasperan/embeddinggemma-code-search")
query_emb = model.encode(["where do we retry failed tool calls?"])
doc_embs = model.encode(code_chunks)
# Matryoshka: truncate + re-normalize for a smaller index
import numpy as np
q128 = query_emb[:, :128]
q128 = q128 / np.linalg.norm(q128, axis=1, keepdims=True)
License
EmbeddingGemma is provided under and subject to the Gemma Terms of Use. This fine-tune inherits those terms.